 archived/Amazon_JumpStart_Semantic_Segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d247be82",
   "metadata": {},
   "source": [
    "# Introduction to JumpStart - Semantic Segmentation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c4ffae36",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c09f5a1",
   "metadata": {},
   "source": [
    "---\n",
    "Welcome to Amazon [SageMaker JumpStart](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-jumpstart.html)! You can use JumpStart to solve many Machine Learning tasks through one-click in SageMaker Studio, or through [SageMaker JumpStart API](https://sagemaker.readthedocs.io/en/stable/doc_utils/jumpstart.html). In this demo notebook, we demonstrate how to use the JumpStart API for Semantic Segmentation. Semantic segmentation is the task of detecting and delineating each distinct object of interest appearing in an image. It is a fine-grained, pixel-level approach to developing computer vision applications. It tags every pixel in an image with a class label from a predefined set of classes. It differs from Instance Segmentation in the following: Semantic Segmentation treats multiple objects of the same class as a single entity whereas Instance Segmentation treats multiple objects of the same class as distinct individual instances.\n",
    "\n",
    "\n",
    "In this notebook, we demonstrate two use cases of semantic segmentation models: \n",
    "\n",
    "* How to use pre-trained Semantic Segmentation models for inference.\n",
    "* How to use JumpStart transfer learning algorithm to finetune a Semantic Segmentation model on a custom dataset.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d299dd5",
   "metadata": {},
   "source": [
    "1. [Set Up](#1.-Set-Up)\n",
    "2. [Select a pre-trained model](#2.-Select-a-pre-trained-model)\n",
    "3. [Run inference on the pre-trained model](#3.-Run-inference-on-the-pre-trained-model)\n",
    "    * [Retrieve JumpStart Artifacts & Deploy an Endpoint](#3.1.-Retrieve-JumpStart-Artifacts-&-Deploy-an-Endpoint)\n",
    "    * [Download an example image for inference](#3.2.-Download-an-example-image-for-inference)\n",
    "    * [Query endpoint and parse response](#3.3.-Query-endpoint-and-parse-response)\n",
    "    * [Display model predictions](#3.4.-Display-model-predictions)\n",
    "    * [Clean up the endpoint](#3.5.-Clean-up-the-endpoint)\n",
    "4. [Fine-tune the pre-trained model on a custom dataset](#4.-Fine-tune-the-pre-trained-model-on-a-custom-dataset)\n",
    "    * [Retrieve Training Artifacts](#4.1.-Retrieve-Training-Artifacts)\n",
    "    * [Set Training parameters](#4.2.-Set-Training-parameters)\n",
    "    * [Train with Automatic Model Tuning (HPO)](#AMT)\n",
    "    * [Start Training](#4.4.-Start-Training)\n",
    "    * [Deploy and run inference on the fine-tuned model](#4.5.-Deploy-and-run-inference-on-the-fine-tuned-model)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd653481",
   "metadata": {},
   "source": [
    "Note: This notebook was tested on ml.t3.medium instance in Amazon SageMaker Studio with Python 3 (Data Science) kernel and in Amazon SageMaker Notebook instance with conda_python3 kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb03fd0",
   "metadata": {},
   "source": [
    "## 1. Set Up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae299675",
   "metadata": {},
   "source": [
    "---\n",
    "Before executing the notebook, there are some initial steps required for set up. This notebook requires latest version of sagemaker and ipywidgets\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74ee016f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sagemaker ipywidgets --upgrade --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46715465",
   "metadata": {},
   "source": [
    "### Permissions and environment variables\n",
    "\n",
    "---\n",
    "To train and host on Amazon SageMaker, we need to set up and authenticate the use of AWS services. Here, we use the execution role associated with the current notebook as the AWS account role with SageMaker access. It has necessary permissions, including access to your data in S3. \n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a24fc317",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker, boto3, json\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "aws_role = get_execution_role()\n",
    "aws_region = boto3.Session().region_name\n",
    "sess = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26f88c90",
   "metadata": {},
   "source": [
    "### 2. Select a pre-trained model\n",
    "\n",
    "***\n",
    "Here, we download jumpstart model_manifest file from the jumpstart s3 bucket, filter-out all the Semantic Segmentation models and select a model for inference.\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0128005c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import Dropdown\n",
    "\n",
    "# download JumpStart model_manifest file.\n",
    "boto3.client(\"s3\").download_file(\n",
    "    f\"jumpstart-cache-prod-{aws_region}\", \"models_manifest.json\", \"models_manifest.json\"\n",
    ")\n",
    "with open(\"models_manifest.json\", \"rb\") as json_file:\n",
    "    model_list = json.load(json_file)\n",
    "\n",
    "# filter-out all the Semantic Segmentation models from the manifest list.\n",
    "semseg_models = []\n",
    "for model in model_list:\n",
    "    model_id = model[\"model_id\"]\n",
    "    if \"-semseg-\" in model_id and model_id not in semseg_models:\n",
    "        semseg_models.append(model_id)\n",
    "\n",
    "print(f\"\\033[38;5;2mChose a model: \\033[0;0m\\n\")\n",
    "\n",
    "# display the model-ids in a dropdown to select a model for inference.\n",
    "model_dropdown = Dropdown(\n",
    "    options=semseg_models,\n",
    "    value=\"mxnet-semseg-fcn-resnet50-ade\",\n",
    "    description=\"\",\n",
    "    style={\"description_width\": \"initial\"},\n",
    "    layout={\"width\": \"max-content\"},\n",
    ")\n",
    "display(model_dropdown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "793de402",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_version=\"*\" fetches the latest version of the model\n",
    "model_id, model_version = model_dropdown.value, \"*\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9ae6f12",
   "metadata": {},
   "source": [
    "## 3. Run inference on the pre-trained model\n",
    "\n",
    "***\n",
    "\n",
    "Using JumpStart, we can perform inference on the pre-trained model, even without fine-tuning it first on a new dataset. \n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11f850da",
   "metadata": {},
   "source": [
    "### 3.1. Retrieve JumpStart Artifacts & Deploy an Endpoint\n",
    "\n",
    "***\n",
    "We retrieve the `deploy_image_uri`, `deploy_source_uri`, and `base_model_uri` for the pre-trained model. To host the pre-trained base-model, we create an instance of [`sagemaker.model.Model`](https://sagemaker.readthedocs.io/en/stable/api/inference/model.html) and deploy it.\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26e79aa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import image_uris, model_uris, script_uris, hyperparameters\n",
    "from sagemaker.model import Model\n",
    "from sagemaker.predictor import Predictor\n",
    "from sagemaker.utils import name_from_base\n",
    "\n",
    "\n",
    "endpoint_name = name_from_base(f\"jumpstart-example-infer-{model_id}\")\n",
    "\n",
    "inference_instance_type = \"ml.p2.xlarge\"\n",
    "\n",
    "# Retrieve the inference docker container uri\n",
    "deploy_image_uri = image_uris.retrieve(\n",
    "    region=None,\n",
    "    framework=None,  # automatically inferred from model_id\n",
    "    image_scope=\"inference\",\n",
    "    model_id=model_id,\n",
    "    model_version=model_version,\n",
    "    instance_type=inference_instance_type,\n",
    ")\n",
    "\n",
    "# Retrieve the inference script uri. This includes scripts for model loading, inference handling etc.\n",
    "deploy_source_uri = script_uris.retrieve(\n",
    "    model_id=model_id, model_version=model_version, script_scope=\"inference\"\n",
    ")\n",
    "\n",
    "\n",
    "# Retrieve the base model uri\n",
    "base_model_uri = model_uris.retrieve(\n",
    "    model_id=model_id, model_version=model_version, model_scope=\"inference\"\n",
    ")\n",
    "\n",
    "\n",
    "# Create the SageMaker model instance\n",
    "model = Model(\n",
    "    image_uri=deploy_image_uri,\n",
    "    source_dir=deploy_source_uri,\n",
    "    model_data=base_model_uri,\n",
    "    entry_point=\"inference.py\",  # entry point file in source_dir and present in deploy_source_uri\n",
    "    role=aws_role,\n",
    "    predictor_cls=Predictor,\n",
    "    name=endpoint_name,\n",
    ")\n",
    "\n",
    "# deploy the Model. Note that we need to pass Predictor class when we deploy model through Model class,\n",
    "# for being able to run inference through the sagemaker API.\n",
    "base_model_predictor = model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=inference_instance_type,\n",
    "    predictor_cls=Predictor,\n",
    "    endpoint_name=endpoint_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e8e1828",
   "metadata": {},
   "source": [
    "### 3.2. Download an example image for inference\n",
    "---\n",
    "We download an example image from the JumpStart assets S3 bucket.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86d51b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "jumpstart_assets_s3_bucket = f\"jumpstart-cache-prod-{aws_region}\"\n",
    "pedestrian_img_key_prefix = \"inference-notebook-assets\"\n",
    "pedestrian_img = \"img_pedestrian.png\"\n",
    "\n",
    "boto3.client(\"s3\").download_file(\n",
    "    jumpstart_assets_s3_bucket, f\"{pedestrian_img_key_prefix}/{pedestrian_img}\", pedestrian_img\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e309dc92",
   "metadata": {},
   "source": [
    "### 3.3. Query endpoint and parse response\n",
    "\n",
    "---\n",
    "Input to the endpoint is a single image in binary format. Response of the endpoint is a predicted label for each pixel in the image. \n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7daac4c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "def query(model_predictor, image_file_name):\n",
    "    \"\"\"Query the model predictor.\"\"\"\n",
    "\n",
    "    with open(image_file_name, \"rb\") as file:\n",
    "        input_img_rb = file.read()\n",
    "\n",
    "    query_response = model_predictor.predict(\n",
    "        input_img_rb,\n",
    "        {\n",
    "            \"ContentType\": \"application/x-image\",\n",
    "            \"Accept\": \"application/json;verbose\",\n",
    "        },\n",
    "    )\n",
    "    return query_response\n",
    "\n",
    "\n",
    "def parse_response(query_response):\n",
    "    \"\"\"Parse response and return predictions as well as the set of all labels and object labels present in the image.\"\"\"\n",
    "    response_dict = json.loads(query_response)\n",
    "    return response_dict[\"predictions\"], response_dict[\"labels\"], response_dict[\"image_labels\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "341f1a88",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_response = query(base_model_predictor, pedestrian_img)\n",
    "\n",
    "predictions, labels, image_labels = parse_response(query_response)\n",
    "print(\"Objects present in the picture:\", image_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1b870ab",
   "metadata": {},
   "source": [
    "### 3.4. Display model predictions\n",
    "---\n",
    "Next, we display the bounding boxes overlaid on the original image. To get color palette for visualization, we borrow the VOC palette implementation from [GluonCV](https://cv.gluon.ai/_modules/gluoncv/utils/viz/segmentation.html#get_color_pallete)  \n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47aed71a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.image as mpimg\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "def getvocpalette(num_cls):\n",
    "    \"\"\"Get a color palette.\"\"\"\n",
    "\n",
    "    n = num_cls\n",
    "    palette = [0] * (n * 3)\n",
    "    for j in range(0, n):\n",
    "        lab = j\n",
    "        palette[j * 3 + 0] = 0\n",
    "        palette[j * 3 + 1] = 0\n",
    "        palette[j * 3 + 2] = 0\n",
    "        i = 0\n",
    "        while lab > 0:\n",
    "            palette[j * 3 + 0] |= ((lab >> 0) & 1) << (7 - i)\n",
    "            palette[j * 3 + 1] |= ((lab >> 1) & 1) << (7 - i)\n",
    "            palette[j * 3 + 2] |= ((lab >> 2) & 1) << (7 - i)\n",
    "            i = i + 1\n",
    "            lab >>= 3\n",
    "    return palette\n",
    "\n",
    "\n",
    "def display_predictions(predictions):\n",
    "    \"\"\"Display predictions with each pixel subsituted by the color of the corresponding label.\"\"\"\n",
    "\n",
    "    palette = getvocpalette(256)\n",
    "    npimg = np.array(predictions)\n",
    "    npimg[npimg == -1] = 255\n",
    "    mask = Image.fromarray(npimg.astype(\"uint8\"))\n",
    "\n",
    "    mask.putpalette(palette)\n",
    "    plt.imshow(mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0f9f6d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_predictions(predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bc77ffb",
   "metadata": {},
   "source": [
    "### 3.5. Clean up the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b8ce5cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete the SageMaker endpoint\n",
    "base_model_predictor.delete_model()\n",
    "base_model_predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e722b5eb",
   "metadata": {},
   "source": [
    "## 4. Fine-tune the pre-trained model on a custom dataset\n",
    "\n",
    "---\n",
    "Previously, we saw how to run inference on a pre-trained model. Next, we discuss how a model can be finetuned to a custom dataset with any number of classes. \n",
    "\n",
    "The model available for fine-tuning build a fully convolutional network (FCN) \"head\" on top of the base network. The fine-tuning step fine-tunes the FCNHead while keeping the parameters of the rest of the model frozen, and returns the fine-tuned model. The objective is to minimize per-pixel Softmax Cross Entropy Loss to train FCN. The model returned by fine-tuning can be further deployed for inference. Below are the instructions for how the training data should be formatted for input to the model.\n",
    "\n",
    "* Input: A directory with sub-directories images, masks and a file class_label_to_prediction_index.json.\n",
    "\n",
    "* Output: A trained model that can be deployed for inference.\n",
    "    * A label mapping file is saved along with the trained model file on the s3 bucket.\n",
    "\n",
    "The input directory should look like below if the training data contains two images. The names of .png files can be anything.\n",
    "\n",
    "\n",
    "* input_directory\n",
    "    * images\n",
    "        * abc.png\n",
    "        * def.png\n",
    "    * masks\n",
    "        * abc.png\n",
    "        * def.png\n",
    "    * class_label_to_prediction_index.json\n",
    "\n",
    "The mask files should have class label information for each pixel.\n",
    "\n",
    "\n",
    "We provide pennfudanped dataset as a default dataset for fine-tuning the model.\n",
    "PennFudanPed comprises images of pedestrians. The dataset has been downloaded from [here](https://www.cis.upenn.edu/~jshi/ped_html/#pub1). \n",
    "\n",
    "Citation:\n",
    "<sub><sup>\n",
    "@ONLINE {pennfudanped,\n",
    "author = \"Liming Wang1, Jianbo Shi2, Gang Song2, and I-fan Shen1\",\n",
    "title = \"Penn-Fudan Database for Pedestrian Detection and Segmentation\",\n",
    "year = \"2007\",\n",
    "url = \"https://www.cis.upenn.edu/~jshi/ped_html/\" }\n",
    "</sup></sub>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23d0080f",
   "metadata": {},
   "source": [
    "### 4.1. Retrieve Training Artifacts\n",
    "\n",
    "Here, we retrieve the training docker container, the training algorithm source, and the pre-trained base model. Note that model_version=\"*\" fetches the latest model.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ee92b70",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import image_uris, model_uris, script_uris\n",
    "\n",
    "train_scope = \"training\"\n",
    "training_instance_type = \"ml.p3.2xlarge\"\n",
    "\n",
    "# Retrieve the docker image\n",
    "train_image_uri = image_uris.retrieve(\n",
    "    region=None,\n",
    "    framework=None,  # automatically inferred from model_id\n",
    "    model_id=model_id,\n",
    "    model_version=model_version,\n",
    "    image_scope=train_scope,\n",
    "    instance_type=training_instance_type,\n",
    ")\n",
    "\n",
    "# Retrieve the training script. This contains all the necessary files including data processing, model training etc.\n",
    "train_source_uri = script_uris.retrieve(\n",
    "    model_id=model_id, model_version=model_version, script_scope=train_scope\n",
    ")\n",
    "# Retrieve the pre-trained model tarball to further fine-tune\n",
    "train_model_uri = model_uris.retrieve(\n",
    "    model_id=model_id, model_version=model_version, model_scope=train_scope\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1a60f66",
   "metadata": {},
   "source": [
    "### 4.2. Set Training parameters\n",
    "\n",
    "---\n",
    "Now that we are done with all the set up that is needed, we are ready to train our Semantic Segmentation model. To begin, let us create a [``sageMaker.estimator.Estimator``](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html) object. This estimator will launch the training job. \n",
    "\n",
    "There are two kinds of parameters that need to be set for training. The first one are the parameters for the training job. These include: (i) Training data path. This is S3 folder in which the input data is stored, (ii) Output path: This the s3 folder in which the training output is stored. (iii) Training instance type: This indicates the type of machine on which to run the training. Typically, we use GPU instances for these training. We defined the training instance type above to fetch the correct train_image_uri. \n",
    "\n",
    "The second set of parameters are algorithm specific training hyper-parameters. \n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6700625",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample training data is available in this bucket\n",
    "training_data_bucket = f\"jumpstart-cache-prod-{aws_region}\"\n",
    "training_data_prefix = \"training-datasets/PennFudanPed_SemSeg/\"\n",
    "\n",
    "training_dataset_s3_path = f\"s3://{training_data_bucket}/{training_data_prefix}\"\n",
    "\n",
    "output_bucket = sess.default_bucket()\n",
    "output_prefix = \"jumpstart-example-semseg-training\"\n",
    "\n",
    "s3_output_location = f\"s3://{output_bucket}/{output_prefix}/output\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79e8344c",
   "metadata": {},
   "source": [
    "---\n",
    "For algorithm specific hyper-parameters, we start by fetching python dictionary of the training hyper-parameters that the algorithm accepts with their default values. This can then be overridden to custom values.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e79b04ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import hyperparameters\n",
    "\n",
    "# Retrieve the default hyper-parameters for fine-tuning the model\n",
    "hyperparameters = hyperparameters.retrieve_default(model_id=model_id, model_version=model_version)\n",
    "\n",
    "# [Optional] Override default hyperparameters with custom values\n",
    "hyperparameters[\"epochs\"] = \"4\"\n",
    "print(hyperparameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "482f08e8",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 4.3. Train with Automatic Model Tuning ([HPO](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning.html)) <a id='AMT'></a>\n",
    "***\n",
    "Amazon SageMaker automatic model tuning, also known as hyperparameter tuning, finds the best version of a model by running many training jobs on your dataset using the algorithm and ranges of hyperparameters that you specify. It then chooses the hyperparameter values that result in a model that performs the best, as measured by a metric that you choose. We will use a [HyperparameterTuner](https://sagemaker.readthedocs.io/en/stable/api/training/tuner.html) object to interact with Amazon SageMaker hyperparameter tuning APIs.\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04d7c91c",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from sagemaker.tuner import ContinuousParameter\n",
    "\n",
    "# Use AMT for tuning and selecting the best model\n",
    "use_amt = True\n",
    "\n",
    "# Define objective metric per framework, based on which the best model will be selected.\n",
    "metric_definitions_per_model = {\n",
    "    \"mxnet\": {\n",
    "        \"metrics\": [{\"Name\": \"val_loss\", \"Regex\": \"validation loss=([0-9\\\\.]+)\"}],\n",
    "        \"type\": \"Minimize\",\n",
    "    }\n",
    "}\n",
    "\n",
    "# You can select from the hyperparameters supported by the model, and configure ranges of values to be searched for training the optimal model.(https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-define-ranges.html)\n",
    "hyperparameter_ranges = {\n",
    "    \"adam-learning-rate\": ContinuousParameter(0.0001, 0.1, scaling_type=\"Logarithmic\")\n",
    "}\n",
    "\n",
    "# Increase the total number of training jobs run by AMT, for increased accuracy (and training time).\n",
    "max_jobs = 6\n",
    "# Change parallel training jobs run by AMT to reduce total training time, constrained by your account limits.\n",
    "# if max_jobs=max_parallel_jobs then Bayesian search turns to Random.\n",
    "max_parallel_jobs = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e15d0ada",
   "metadata": {},
   "source": [
    "### 4.4. Start Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd456ad1",
   "metadata": {},
   "source": [
    "---\n",
    "We start by creating the estimator object with all the required assets and then launch the training job.  It takes less than 30 mins on the default dataset.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b90a3d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.estimator import Estimator\n",
    "from sagemaker.utils import name_from_base\n",
    "from sagemaker.tuner import HyperparameterTuner\n",
    "\n",
    "training_job_name = name_from_base(f\"jumpstart-example-{model_id}-transfer-learning\")\n",
    "\n",
    "# Create SageMaker Estimator instance\n",
    "semseg_estimator = Estimator(\n",
    "    role=aws_role,\n",
    "    image_uri=train_image_uri,\n",
    "    source_dir=train_source_uri,\n",
    "    model_uri=train_model_uri,\n",
    "    entry_point=\"transfer_learning.py\",  # Entry-point file in source_dir and present in train_source_uri.\n",
    "    instance_count=1,\n",
    "    instance_type=training_instance_type,\n",
    "    max_run=360000,\n",
    "    hyperparameters=hyperparameters,\n",
    "    output_path=s3_output_location,\n",
    "    base_job_name=training_job_name,\n",
    ")\n",
    "\n",
    "if use_amt:\n",
    "    metric_definitions = next(\n",
    "        value for key, value in metric_definitions_per_model.items() if model_id.startswith(key)\n",
    "    )\n",
    "\n",
    "    hp_tuner = HyperparameterTuner(\n",
    "        semseg_estimator,\n",
    "        metric_definitions[\"metrics\"][0][\"Name\"],\n",
    "        hyperparameter_ranges,\n",
    "        metric_definitions[\"metrics\"],\n",
    "        max_jobs=max_jobs,\n",
    "        max_parallel_jobs=max_parallel_jobs,\n",
    "        objective_type=metric_definitions[\"type\"],\n",
    "        base_tuning_job_name=training_job_name,\n",
    "    )\n",
    "\n",
    "    # Launch a SageMaker Tuning job to search for the best hyperparameters\n",
    "    hp_tuner.fit({\"training\": training_dataset_s3_path})\n",
    "else:\n",
    "    # Launch a SageMaker Training job by passing s3 path of the training data\n",
    "    semseg_estimator.fit({\"training\": training_dataset_s3_path}, logs=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82017a86",
   "metadata": {},
   "source": [
    "### 4.5. Deploy and run inference on the fine-tuned model\n",
    "\n",
    "---\n",
    "\n",
    "A trained model does nothing on its own. We now want to use the model to perform inference. We follow the same steps as in [3. Run inference on the pre-trained model](#3.-Run-inference-on-the-pre-trained-model). We start by retrieving the jumpstart artifacts for deploying an endpoint. However, instead of base_predictor, we  deploy the `semseg_estimator` that we fine-tuned.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71a6ff29",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_instance_type = \"ml.p2.xlarge\"\n",
    "\n",
    "# Retrieve the inference docker container uri\n",
    "deploy_image_uri = image_uris.retrieve(\n",
    "    region=None,\n",
    "    framework=None,  # automatically inferred from model_id\n",
    "    image_scope=\"inference\",\n",
    "    model_id=model_id,\n",
    "    model_version=model_version,\n",
    "    instance_type=inference_instance_type,\n",
    ")\n",
    "# Retrieve the inference script uri. This includes scripts for model loading, inference handling etc.\n",
    "deploy_source_uri = script_uris.retrieve(\n",
    "    model_id=model_id, model_version=model_version, script_scope=\"inference\"\n",
    ")\n",
    "\n",
    "endpoint_name = name_from_base(f\"jumpstart-example-FT-{model_id}-\")\n",
    "\n",
    "# Use the estimator from the previous step to deploy to a SageMaker endpoint\n",
    "finetuned_predictor = (hp_tuner if use_amt else semseg_estimator).deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=inference_instance_type,\n",
    "    entry_point=\"inference.py\",  # entry point file in source_dir and present in deploy_source_uri\n",
    "    image_uri=deploy_image_uri,\n",
    "    source_dir=deploy_source_uri,\n",
    "    endpoint_name=endpoint_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bce41b9",
   "metadata": {},
   "source": [
    "---\n",
    "Next, we download an example pedestrian image from the S3 bucket for inference.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67cf5dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "jumpstart_assets_bucket = f\"jumpstart-cache-prod-{aws_region}\"\n",
    "pedestrian_image_key = \"training-datasets/PennFudanPed_SemSeg/images\"\n",
    "pedestrian_image_file_name = \"FudanPed00001.png\"\n",
    "\n",
    "boto3.client(\"s3\").download_file(\n",
    "    jumpstart_assets_bucket,\n",
    "    f\"{pedestrian_image_key}/{pedestrian_image_file_name}\",\n",
    "    pedestrian_image_file_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86329d32",
   "metadata": {},
   "source": [
    "---\n",
    "Next, we query the finetuned model, parse the response and display the predictions. Functions for these are implemented in sections [3.3. Query endpoint and parse response](#3.3.-Query-endpoint-and-parse-response) and [3.4. Display model predictions](#3.4.-Display-model-predictions)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af7b7c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_response = query(finetuned_predictor, pedestrian_img)\n",
    "predictions, labels, image_labels = parse_response(query_response)\n",
    "display_predictions(predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "214f5d8c",
   "metadata": {},
   "source": [
    "---\n",
    "Next, we delete the endpoint corresponding to the finetuned model.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c71283c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete the SageMaker endpoint\n",
    "finetuned_predictor.delete_model()\n",
    "finetuned_predictor.delete_endpoint()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e37ebe4e",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/introduction_to_amazon_algorithms|jumpstart_semantic_segmentation|Amazon_JumpStart_Semantic_Segmentation.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science 3.0)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/sagemaker-data-science-310-v1"
  },
  "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}