1_model_training/1_train_model.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4e1cd147",
   "metadata": {},
   "source": [
    "# Tensorflow Object Detection API and AWS Sagemaker"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "85592c17",
   "metadata": {},
   "source": [
    "In this notebook, you will train and evaluate different models using the [Tensorflow Object Detection API](https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/) and [AWS Sagemaker](https://aws.amazon.com/sagemaker/). \n",
    "\n",
    "If you ever feel stuck, you can refer to this [tutorial](https://aws.amazon.com/blogs/machine-learning/training-and-deploying-models-using-tensorflow-2-with-the-object-detection-api-on-amazon-sagemaker/).\n",
    "\n",
    "## Dataset\n",
    "\n",
    "We are using the [Waymo Open Dataset](https://waymo.com/open/) for this project. The dataset has already been exported using the tfrecords format. The files have been created following the format described [here](https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/training.html#create-tensorflow-records). You can find data stored on [AWS S3](https://aws.amazon.com/s3/), AWS Object Storage. The images are saved with a resolution of 640x640."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc1d114",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install tensorflow_io sagemaker -U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96f55350",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sagemaker\n",
    "from sagemaker.estimator import Estimator\n",
    "from framework import CustomFramework"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ccde6fd1",
   "metadata": {},
   "source": [
    "Save the IAM role in a variable called `role`. This would be useful when training the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab6b13b",
   "metadata": {},
   "outputs": [],
   "source": [
    "role = sagemaker.get_execution_role()\n",
    "print(role)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae64e29f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The train and val paths below are public S3 buckets created by Udacity for this project\n",
    "inputs = {'train': 's3://cd2688-object-detection-tf2/train/', \n",
    "          'val': 's3://cd2688-object-detection-tf2/val/'} \n",
    "\n",
    "# Insert path of a folder in your personal S3 bucket to store tensorboard logs.\n",
    "tensorboard_s3_prefix = 's3://object-detection-project/logs/'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fc16a825",
   "metadata": {},
   "source": [
    "## Container\n",
    "\n",
    "To train the model, you will first need to build a [docker](https://www.docker.com/) container with all the dependencies required by the TF Object Detection API. The code below does the following:\n",
    "* clone the Tensorflow models repository\n",
    "* get the exporter and training scripts from the repository\n",
    "* build the docker image and push it \n",
    "* print the container name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad5ac8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# clone the repo and get the scripts\n",
    "git clone https://github.com/tensorflow/models.git docker/models\n",
    "\n",
    "# get model_main and exporter_main files from TF2 Object Detection GitHub repository\n",
    "cp docker/models/research/object_detection/exporter_main_v2.py source_dir \n",
    "cp docker/models/research/object_detection/model_main_tf2.py source_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2dab3f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# build and push the docker image. This code can be commented out after being run once.\n",
    "# This will take around 10 mins.\n",
    "image_name = 'tf2-object-detection'\n",
    "!sh ./docker/build_and_push.sh $image_name"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e62b3562",
   "metadata": {},
   "source": [
    "To verify that the image was correctly pushed to the [Elastic Container Registry](https://aws.amazon.com/ecr/), you can look at it in the AWS webapp. For example, below you can see that three different images have been pushed to ECR. You should only see one, called `tf2-object-detection`.\n",
    "![ECR Example](../data/example_ecr.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0310b6a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# display the container name\n",
    "with open (os.path.join('docker', 'ecr_image_fullname.txt'), 'r') as f:\n",
    "    container = f.readlines()[0][:-1]\n",
    "\n",
    "print(container)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "13b2a754",
   "metadata": {},
   "source": [
    "## Pre-trained model from model zoo\n",
    "\n",
    "As often, we are not training from scratch and we will be using a pretrained model from the TF Object Detection model zoo. You can find pretrained checkpoints [here](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md). Because your time is limited for this project, we recommend to only experiment with the following models:\n",
    "* SSD MobileNet V2 FPNLite 640x640\t\n",
    "* SSD ResNet50 V1 FPN 640x640 (RetinaNet50)\t\n",
    "* Faster R-CNN ResNet50 V1 640x640\t\n",
    "* EfficientDet D1 640x640\t\n",
    "* Faster R-CNN ResNet152 V1 640x640\t\n",
    "\n",
    "In the code below, the EfficientDet D1 model is downloaded and extracted. This code should be adjusted if you were to experiment with other architectures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c4b1d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "mkdir /tmp/checkpoint\n",
    "mkdir source_dir/checkpoint\n",
    "wget -O /tmp/efficientdet.tar.gz http://download.tensorflow.org/models/object_detection/tf2/20200711/efficientdet_d1_coco17_tpu-32.tar.gz\n",
    "tar -zxvf /tmp/efficientdet.tar.gz --strip-components 2 --directory source_dir/checkpoint efficientdet_d1_coco17_tpu-32/checkpoint"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a1e04a98",
   "metadata": {},
   "source": [
    "## Edit pipeline.config file\n",
    "\n",
    "The [`pipeline.config`](source_dir/pipeline.config) in the `source_dir` folder should be updated when you experiment with different models. The different config files are available [here](https://github.com/tensorflow/models/tree/master/research/object_detection/configs/tf2).\n",
    "\n",
    ">Note: The provided `pipeline.config` file works well with the `EfficientDet` model. You would need to modify it when working with other models."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "47483545",
   "metadata": {},
   "source": [
    "## Launch Training Job\n",
    "\n",
    "Now that we have a dataset, a docker image and some pretrained model weights, we can launch the training job. To do so, we create a [Sagemaker Framework](https://sagemaker.readthedocs.io/en/stable/frameworks/index.html), where we indicate the container name, name of the config file, number of training steps etc.\n",
    "\n",
    "The `run_training.sh` script does the following:\n",
    "* train the model for `num_train_steps` \n",
    "* evaluate over the val dataset\n",
    "* export the model\n",
    "\n",
    "Different metrics will be displayed during the evaluation phase, including the mean average precision. These metrics can be used to quantify your model performances and compare over the different iterations.\n",
    "\n",
    "You can also monitor the training progress by navigating to **Training -> Training Jobs** from the Amazon Sagemaker dashboard in the Web UI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c7175cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorboard_output_config = sagemaker.debugger.TensorBoardOutputConfig(\n",
    "    s3_output_path=tensorboard_s3_prefix,\n",
    "    container_local_output_path='/opt/training/'\n",
    ")\n",
    "\n",
    "estimator = CustomFramework(\n",
    "    role=role,\n",
    "    image_uri=container,\n",
    "    entry_point='run_training.sh',\n",
    "    source_dir='source_dir/',\n",
    "    hyperparameters={\n",
    "        \"model_dir\": \"/opt/training\",        \n",
    "        \"pipeline_config_path\": \"pipeline.config\",\n",
    "        \"num_train_steps\": \"2000\",    \n",
    "        \"sample_1_of_n_eval_examples\": \"1\"\n",
    "    },\n",
    "    instance_count=1,\n",
    "    instance_type='ml.g5.xlarge',\n",
    "    tensorboard_output_config=tensorboard_output_config,\n",
    "    disable_profiler=True,\n",
    "    base_job_name='tf2-object-detection'\n",
    ")\n",
    "\n",
    "estimator.fit(inputs)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "84545881",
   "metadata": {},
   "source": [
    "You should be able to see your model training in the AWS webapp as shown below:\n",
    "![ECR Example](../data/example_trainings.png)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f9844f25",
   "metadata": {},
   "source": [
    "## Improve on the initial model\n",
    "\n",
    "Most likely, this initial experiment did not yield optimal results. However, you can make multiple changes to the `pipeline.config` file to improve this model. One obvious change consists in improving the data augmentation strategy. The [`preprocessor.proto`](https://github.com/tensorflow/models/blob/master/research/object_detection/protos/preprocessor.proto) file contains the different data augmentation method available in the Tf Object Detection API. Justify your choices of augmentations in the write-up.\n",
    "\n",
    "Keep in mind that the following are also available:\n",
    "* experiment with the optimizer: type of optimizer, learning rate, scheduler etc\n",
    "* experiment with the architecture. The Tf Object Detection API model zoo offers many architectures. Keep in mind that the pipeline.config file is unique for each architecture and you will have to edit it.\n",
    "* visualize results on the test frames using the `2_deploy_model` notebook available in this repository.\n",
    "\n",
    "In the cell below, write down all the different approaches you have experimented with, why you have chosen them and what you would have done if you had more time and resources. Justify your choices using the tensorboard visualizations (take screenshots and insert them in your write-up), the metrics on the evaluation set and the generated animation you have created with [this tool](../2_run_inference/2_deploy_model.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17284a60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# your write-up goes here."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7 (default, Sep 16 2021, 08:50:36) \n[Clang 10.0.0 ]"
  },
  "vscode": {
   "interpreter": {
    "hash": "40d3a090f54c6569ab1632332b64b2c03c39dcf918b08424e98f38b5ae0af88f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}