{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using SageMaker Debugger and SageMaker Experiments for iterative model pruning\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "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://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-west-2/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This notebook demonstrates how we can use [SageMaker Debugger](https://docs.aws.amazon.com/sagemaker/latest/dg/train-debugger.html) and [SageMaker Experiments](https://docs.aws.amazon.com/sagemaker/latest/dg/experiments.html) to perform iterative model pruning. Let's start first with a quick introduction into model pruning.\n",
    "\n",
    "State-of-the-art deep learning models consist of millions of parameters and are trained on very large datasets. For transfer learning we take a pre-trained model and fine-tune it on a new and typically much smaller dataset. The new dataset may even consist of different classes, so the model is basically learning a new task. This process allows us to quickly achieve state-of-the-artresults without having to design and train our own model from scratch. However, it may happen that a much smaller and simpler model would also perform well on our dataset. With model pruning we identify the importance of weights during training and remove the weights that are contributing very little to the learning process. We can do this in an iterative way where we remove a small percentage of weights in each iteration. Removing means to eliminate the entries in the tensor so its size shrinks.\n",
    "\n",
    "We use SageMaker Debugger to get weights, activation outputs and gradients during training. These tensors are used to compute the importance of weights. We will use SageMaker Experiments to keep track of each pruning iteration: if we prune too much we may degrade model accuracy, so we will monitor number of parameters versus validation accuracy. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pip\n",
    "import sys\n",
    "\n",
    "def import_or_install(package):\n",
    "    try:\n",
    "        __import__(package)\n",
    "    except ImportError:\n",
    "        !{sys.executable} -m pip install {package}\n",
    "        \n",
    "required_packages=['smdebug', 'sagemaker-experiments']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get training dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we get the [Caltech-101](https://paperswithcode.com/dataset/caltech-101) dataset. This dataset consists of 101 image categories. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "import requests\n",
    "import os\n",
    "\n",
    "filename = \"101_ObjectCategories.tar.gz\"\n",
    "data_url = os.path.join(\"https://s3.us-east-2.amazonaws.com/mxnet-public\", filename)\n",
    "\n",
    "r = requests.get(data_url, stream=True)\n",
    "with open(filename, \"wb\") as f:\n",
    "    for chunk in r.iter_content(chunk_size=1024):\n",
    "        if chunk:\n",
    "            f.write(chunk)\n",
    "\n",
    "print(\"Extracting {} ...\".format(filename))\n",
    "tar = tarfile.open(filename, \"r:gz\")\n",
    "tar.extractall(\".\")\n",
    "tar.close()\n",
    "print(\"Data extracted.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And upload it to our SageMaker default bucket:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "\n",
    "\n",
    "def upload_to_s3(path, directory_name, bucket, counter=-1):\n",
    "    print(\"Upload files from\" + path + \" to \" + bucket)\n",
    "    client = boto3.client(\"s3\")\n",
    "\n",
    "    for path, subdirs, files in os.walk(path):\n",
    "        path = path.replace(\"\\\\\", \"/\")\n",
    "        print(path)\n",
    "        for file in files[0:counter]:\n",
    "            client.upload_file(\n",
    "                os.path.join(path, file),\n",
    "                bucket,\n",
    "                directory_name + \"/\" + path.split(\"/\")[-1] + \"/\" + file,\n",
    "            )\n",
    "\n",
    "\n",
    "boto_session = boto3.Session()\n",
    "sagemaker_session = sagemaker.Session(boto_session=boto_session)\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "\n",
    "upload_to_s3(\"101_ObjectCategories\", directory_name=\"101_ObjectCategories_train\", bucket=bucket)\n",
    "\n",
    "# we will compute saliency maps for all images in the test dataset, so we will only upload 4 images\n",
    "upload_to_s3(\n",
    "    \"101_ObjectCategories_test\",\n",
    "    directory_name=\"101_ObjectCategories_test\",\n",
    "    bucket=bucket,\n",
    "    counter=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load and save ResNet model\n",
    "\n",
    "First we load a pre-trained [ResNet](https://arxiv.org/abs/1512.03385) model from PyTorch model zoo. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import models\n",
    "from torch import nn\n",
    "\n",
    "model = models.resnet18(pretrained=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look on the model architecture:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see above, the last Linear layer outputs 1000 values, which is the number of classes the model has originally been trained on. Here, we will fine-tune the model on the Caltech101 dataset: as it has only 101 classes, we need to set the number of output classes to 101."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nfeatures = model.fc.in_features\n",
    "model.fc = torch.nn.Linear(nfeatures, 101)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we store the model definition and weights in an output file. \n",
    "\n",
    "**IMPORTANT**: the model file will be used by the training job. To avoid version conflicts, you need to ensure that your notebook is running a Jupyter kernel with PyTorch version 1.6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = {\"model\": model, \"state_dict\": model.state_dict()}\n",
    "\n",
    "torch.save(checkpoint, \"src/model_checkpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code cell creates a SageMaker experiment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from datetime import datetime\n",
    "from smexperiments.experiment import Experiment\n",
    "\n",
    "sagemaker_boto_client = boto3.client(\"sagemaker\")\n",
    "\n",
    "# name of experiment\n",
    "timestep = datetime.now()\n",
    "timestep = timestep.strftime(\"%d-%m-%Y-%H-%M-%S\")\n",
    "experiment_name = timestep + \"resnet-model-pruning-experiment\"\n",
    "\n",
    "# create experiment\n",
    "Experiment.create(\n",
    "    experiment_name=experiment_name,\n",
    "    description=\"Iterative model pruning of ResNet trained on Caltech101\",\n",
    "    sagemaker_boto_client=sagemaker_boto_client,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code cell defines a list of tensor names that be used to compute filter ranks. The lists are defined in the Python script `model_resnet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import model_resnet\n",
    "\n",
    "activation_outputs = model_resnet.activation_outputs\n",
    "gradients = model_resnet.gradients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterative model pruning: step by step\n",
    "\n",
    "Before we jump into the code for running the iterative model pruning we will walk through the code step by step. \n",
    "\n",
    "#### Step 0: Create trial and debugger hook configuration\n",
    "First we create a new trial for each pruning iteration. That allows us to track our training jobs and see which models have the lowest number of parameters and best accuracy. We use the `smexperiments` library to create a trial within our experiment.                       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smexperiments.trial import Trial\n",
    "\n",
    "trial = Trial.create(experiment_name=experiment_name, sagemaker_boto_client=sagemaker_boto_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we define the experiment_config which is a dictionary that will be passed to the SageMaker training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_config = {\n",
    "    \"ExperimentName\": experiment_name,\n",
    "    \"TrialName\": trial.trial_name,\n",
    "    \"TrialComponentDisplayName\": \"Training\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create a debugger hook configuration to define a custom collection of tensors to be emitted. The custom collection contains all weights and biases of the model. It also includes individual layer outputs and their gradients which will be used to compute filter ranks. Tensors are saved every 100th iteration where an iteration represents one forward and backward pass. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.debugger import DebuggerHookConfig, CollectionConfig\n",
    "\n",
    "debugger_hook_config = DebuggerHookConfig(\n",
    "    collection_configs=[\n",
    "        CollectionConfig(\n",
    "            name=\"custom_collection\",\n",
    "            parameters={\n",
    "                \"include_regex\": \".*relu|.*weight|.*bias|.*running_mean|.*running_var|.*CrossEntropyLoss\",\n",
    "                \"save_interval\": \"100\",\n",
    "            },\n",
    "        )\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 1: Start training job\n",
    "Now we define the SageMaker PyTorch Estimator. We will train the model on an `ml.p2.xlarge` instance. The model definition plus training code is defined in the entry_point file `train.py`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "estimator = PyTorch(\n",
    "    role=sagemaker.get_execution_role(),\n",
    "    instance_count=1,\n",
    "    instance_type=\"ml.p3.2xlarge\",\n",
    "    volume_size=400,\n",
    "    source_dir=\"src\",\n",
    "    entry_point=\"train.py\",\n",
    "    framework_version=\"1.12\",\n",
    "    py_version=\"py38\",\n",
    "    metric_definitions=[\n",
    "        {\"Name\": \"train:loss\", \"Regex\": \"loss:(.*?)\"},\n",
    "        {\"Name\": \"eval:acc\", \"Regex\": \"acc:(.*?)\"},\n",
    "    ],\n",
    "    enable_sagemaker_metrics=True,\n",
    "    hyperparameters={\"epochs\": 10},\n",
    "    debugger_hook_config=debugger_hook_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we define the estimator object, we can call `fit` which creates an `ml.p2.xlarge` instance on which it starts the training. We pass the experiment_config which associates the training job with a trial and an experiment. If we don't specify an `experiment_config`, the training job will appear in SageMaker Experiments under *Unassigned trial components*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit(\n",
    "    inputs={\n",
    "        \"train\": \"s3://{}/101_ObjectCategories_train\".format(bucket),\n",
    "        \"test\": \"s3://{}/101_ObjectCategories_test\".format(bucket),\n",
    "    },\n",
    "    experiment_config=experiment_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2: Get gradients, weights, biases\n",
    "\n",
    "Once the training job has finished, we will retrieve its tensors, such as gradients, weights and biases. We use the `smdebug` library which provides functions to read and filter tensors. First we create a [trial](https://github.com/awslabs/sagemaker-debugger/blob/master/docs/analysis.md#Trial) that is reading the tensors from S3. \n",
    "\n",
    "For clarification: in the context of SageMaker Debugger a trial is an object that lets you query tensors for a given training job. In the context of SageMaker Experiments a trial is part of an experiment; it presents a collection of training steps involved in a single training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from smdebug.trials import create_trial\n",
    "\n",
    "path = estimator.latest_job_debugger_artifacts_path()\n",
    "smdebug_trial = create_trial(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To access tensor values, we only need to call `smdebug_trial.tensor()`. For instance, to get the outputs of the first ReLU activation at step 0, we run `smdebug_trial.tensor('layer4.1.relu_0_output_0').value(0, mode=modes.TRAIN)`. Next we compute a filter rank for the convolutions. \n",
    "\n",
    "Some definitions: a filter is a collection of kernels (one kernel for every single input channel) and a filter produces one feature map (output channel). In the image below the convolution creates 64 feature maps (output channels) and uses a kernel of 5x5. By pruning a filter, an entire feature map will be removed. So in the example image below the number of feature maps (output channels) would shrink to 63 and the number of learnable parameters (weights) would be reduced by 1x5x5.\n",
    "\n",
    "![](images/convolution.png) \n",
    "\n",
    "\n",
    "#### Step 3: Compute filter ranks\n",
    "\n",
    "In this notebook we compute filter ranks as described in the article [\"Pruning Convolutional Neural Networks for Resource Efficient Inference\"](https://arxiv.org/pdf/1611.06440.pdf) We basically identify filters that are less important for the final prediction of the model. The product of weights and gradients can be seen as a measure of importance. The product has the dimension `(batch_size, out_channels, width, height)` and we get the average over `axis=0,2,3` to have a single value (rank) for each filter.\n",
    "\n",
    "In the following code we retrieve activation outputs and gradients and compute the filter rank. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from smdebug import modes\n",
    "\n",
    "\n",
    "def compute_filter_ranks(smdebug_trial, activation_outputs, gradients):\n",
    "    filters = {}\n",
    "    for activation_output_name, gradient_name in zip(activation_outputs, gradients):\n",
    "        for step in smdebug_trial.steps(mode=modes.TRAIN):\n",
    "            activation_output = smdebug_trial.tensor(activation_output_name).value(\n",
    "                step, mode=modes.TRAIN\n",
    "            )\n",
    "            gradient = smdebug_trial.tensor(gradient_name).value(step, mode=modes.TRAIN)\n",
    "            rank = activation_output * gradient\n",
    "            rank = np.mean(rank, axis=(0, 2, 3))\n",
    "\n",
    "            if activation_output_name not in filters:\n",
    "                filters[activation_output_name] = 0\n",
    "            filters[activation_output_name] += rank\n",
    "    return filters\n",
    "\n",
    "\n",
    "filters = compute_filter_ranks(smdebug_trial, activation_outputs, gradients)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we normalize the filters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_filter_ranks(filters):\n",
    "    for activation_output_name in filters:\n",
    "        rank = np.abs(filters[activation_output_name])\n",
    "        rank = rank / np.sqrt(np.sum(rank * rank))\n",
    "        filters[activation_output_name] = rank\n",
    "    return filters\n",
    "\n",
    "\n",
    "filters = normalize_filter_ranks(filters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create a list of filters, sort it by rank and retrieve the smallest values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_smallest_filters(filters, n):\n",
    "    filters_list = []\n",
    "    for layer_name in sorted(filters.keys()):\n",
    "        for channel in range(filters[layer_name].shape[0]):\n",
    "            filters_list.append(\n",
    "                (\n",
    "                    layer_name,\n",
    "                    channel,\n",
    "                    filters[layer_name][channel],\n",
    "                )\n",
    "            )\n",
    "\n",
    "    filters_list.sort(key=lambda x: x[2])\n",
    "    filters_list = filters_list[:n]\n",
    "    print(\"The\", n, \"smallest filters\", filters_list)\n",
    "\n",
    "    return filters_list\n",
    "\n",
    "\n",
    "filters_list = get_smallest_filters(filters, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 4 and step 5: Prune low ranking filters and set new weights\n",
    "\n",
    "Next we prune the model, where we remove filters and their corresponding weights. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "step = smdebug_trial.steps(mode=modes.TRAIN)[-1]\n",
    "\n",
    "model = model_resnet.prune(model, filters_list, smdebug_trial, step)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 6: Start next pruning iteration\n",
    "Once we have pruned the model, the new architecture and pruned weights will be saved under src and will be used by the next training job in the next pruning iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save pruned model\n",
    "checkpoint = {\"model\": model, \"state_dict\": model.state_dict()}\n",
    "\n",
    "torch.save(checkpoint, \"src/model_checkpoint\")\n",
    "\n",
    "# clean up\n",
    "del model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Overall workflow\n",
    "The overall workflow looks like the following:\n",
    " ![](images/workflow.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run iterative model pruning\n",
    "\n",
    "After having gone through the code step by step, we are ready to run the full workflow. The following cell runs 1 pruning iteration for a tutorial purpose. Change the range of the for loop to 10 to replicate the same result shown in the [Pruning machine learning models with Amazon SageMaker Debugger and Amazon SageMaker Experiments blog](https://aws.amazon.com/blogs/machine-learning/pruning-machine-learning-models-with-amazon-sagemaker-debugger-and-amazon-sagemaker-experiments/) and the figure below the cell. In each iteration a new SageMaker training job is started, where it emits gradients and activation outputs to Amazon S3. Once the job has finished, filter ranks are computed and the 100 smallest filters are removed.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start iterative pruning\n",
    "for pruning_step in range(1):\n",
    "    # create new trial for this pruning step\n",
    "    smexperiments_trial = Trial.create(\n",
    "        experiment_name=experiment_name, sagemaker_boto_client=sagemaker_boto_client\n",
    "    )\n",
    "    experiment_config[\"TrialName\"] = smexperiments_trial.trial_name\n",
    "\n",
    "    print(\"Created new trial\", smexperiments_trial.trial_name, \"for pruning step\", pruning_step)\n",
    "\n",
    "    # start training job\n",
    "    estimator = PyTorch(\n",
    "        role=sagemaker.get_execution_role(),\n",
    "        instance_count=1,\n",
    "        instance_type=\"ml.p3.2xlarge\",\n",
    "        volume_size=400,\n",
    "        source_dir=\"src\",\n",
    "        entry_point=\"train.py\",\n",
    "        framework_version=\"1.6\",\n",
    "        py_version=\"py3\",\n",
    "        metric_definitions=[\n",
    "            {\"Name\": \"train:loss\", \"Regex\": \"loss:(.*?)\"},\n",
    "            {\"Name\": \"eval:acc\", \"Regex\": \"acc:(.*?)\"},\n",
    "        ],\n",
    "        enable_sagemaker_metrics=True,\n",
    "        hyperparameters={\"epochs\": 10},\n",
    "        debugger_hook_config=debugger_hook_config,\n",
    "    )\n",
    "\n",
    "    # start training job\n",
    "    estimator.fit(\n",
    "        inputs={\n",
    "            \"train\": \"s3://{}/101_ObjectCategories_train\".format(bucket),\n",
    "            \"test\": \"s3://{}/101_ObjectCategories_test\".format(bucket),\n",
    "        },\n",
    "        experiment_config=experiment_config,\n",
    "    )\n",
    "\n",
    "    print(\"Training job\", estimator.latest_training_job.name, \" finished.\")\n",
    "\n",
    "    # read tensors\n",
    "    path = estimator.latest_job_debugger_artifacts_path()\n",
    "    smdebug_trial = create_trial(path)\n",
    "\n",
    "    # compute filter ranks and get 100 smallest filters\n",
    "    filters = compute_filter_ranks(smdebug_trial, activation_outputs, gradients)\n",
    "    filters_normalized = normalize_filter_ranks(filters)\n",
    "    filters_list = get_smallest_filters(filters_normalized, 100)\n",
    "\n",
    "    # load previous model\n",
    "    checkpoint = torch.load(\"src/model_checkpoint\")\n",
    "    model = checkpoint[\"model\"]\n",
    "    model.load_state_dict(checkpoint[\"state_dict\"])\n",
    "\n",
    "    # prune model\n",
    "    step = smdebug_trial.steps(mode=modes.TRAIN)[-1]\n",
    "    model = model_resnet.prune(model, filters_list, smdebug_trial, step)\n",
    "\n",
    "    print(\"Saving pruned model\")\n",
    "\n",
    "    # save pruned model\n",
    "    checkpoint = {\"model\": model, \"state_dict\": model.state_dict()}\n",
    "    torch.save(checkpoint, \"src/model_checkpoint\")\n",
    "\n",
    "    # clean up\n",
    "    del model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the iterative model pruning is running, we can track and visualize our experiment in SageMaker Studio. In our training script we use SageMaker debugger's `save_scalar` method to store the number of parameters in the model and the model accuracy. So we can visualize those in Studio or use the `ExperimentAnalytics` module to read and plot the values directly in the notebook.\n",
    "\n",
    "Initially the model consisted of 11 million parameters. After 11 iterations, the number of parameters was reduced to 270k, while accuracy increased to 91% and then started dropping after 8 pruning iteration.\n",
    "\n",
    "This means that the best accuracy can be reached if the model has a size of about 4 million parameters, while shrinking model size about 3x!\n",
    "\n",
    "![](images/results_resnet.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Additional: run iterative model pruning with custom rule\n",
    "\n",
    "In the previous example, we have seen that accuracy drops when the model has less than 22 million parameters. Clearly, we want to stop our experiment once we reach this point. We can define a custom rule that returns `True` if the accuracy drops by a certain percentage. You can find an example implementation in `custom_rule/check_accuracy.py`. Before we can use the rule we have to define a custom rule configuration:\n",
    "\n",
    "```python\n",
    "\n",
    "from sagemaker.debugger import Rule, CollectionConfig, rule_configs\n",
    "\n",
    "check_accuracy_rule = Rule.custom(\n",
    "    name='CheckAccuracy',\n",
    "    image_uri='759209512951.dkr.ecr.us-west-2.amazonaws.com/sagemaker-debugger-rule-evaluator:latest',\n",
    "    instance_type='ml.c4.xlarge',\n",
    "    volume_size_in_gb=400,\n",
    "    source='custom_rule/check_accuracy.py',\n",
    "    rule_to_invoke='check_accuracy',\n",
    "    rule_parameters={\"previous_accuracy\": \"0.0\", \n",
    "                     \"threshold\": \"0.05\", \n",
    "                     \"predictions\": \"CrossEntropyLoss_0_input_0\", \n",
    "                     \"labels\":\"CrossEntropyLoss_0_input_1\"},\n",
    ")\n",
    "```\n",
    "\n",
    "The rule reads the inputs to the loss function, which are the model predictions and the labels. It computes the accuracy and returns `True` if its value has dropped by more than 5% otherwise `False`. \n",
    "\n",
    "In each pruning iteration, we need to pass the accuracy of the previous training job to the rule, which can be retrieved via the `ExperimentAnalytics` module.\n",
    "\n",
    "```python\n",
    "from sagemaker.analytics import ExperimentAnalytics\n",
    "\n",
    "trial_component_analytics = ExperimentAnalytics(experiment_name=experiment_name)\n",
    "accuracy = trial_component_analytics.dataframe()['scalar/accuracy_EVAL - Max'][0]\n",
    "```\n",
    "And overwrite the value in the rule configuration:\n",
    "\n",
    "```python\n",
    "check_accuracy_rule.rule_parameters[\"previous_accuracy\"] = str(accuracy)\n",
    "```\n",
    "\n",
    "In the PyTorch estimator we need to add the argument `rules = [check_accuracy_rule]`.\n",
    "We can create a CloudWatch alarm and use a Lambda function to stop the training. Detailed instructions can be found [here](https://github.com/awslabs/amazon-sagemaker-examples/tree/master/sagemaker-debugger/tensorflow_action_on_rule). In each iteration we check the job status and if the previous job has been stopped, we exit the loop:\n",
    "\n",
    "```python\n",
    "job_name = estimator.latest_training_job.name\n",
    "client = estimator.sagemaker_session.sagemaker_client\n",
    "description = client.describe_training_job(TrainingJobName=job_name)\n",
    "\n",
    "if description['TrainingJobStatus'] == 'Stopped':\n",
    "      break\n",
    "```\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "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://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-east-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-east-2/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/us-west-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ca-central-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/sa-east-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-west-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-west-2/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-west-3/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-central-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/eu-north-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-southeast-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-southeast-2/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-northeast-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-northeast-2/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://h75twx4l60.execute-api.us-west-2.amazonaws.com/sagemaker-nb/ap-south-1/sagemaker-debugger|pytorch_iterative_model_pruning|iterative_model_pruning_resnet.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.10 Python 3.8 CPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-2:236514542706:image/pytorch-1.10-cpu-py38"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}