{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "924b27fa",
   "metadata": {},
   "source": [
    "# Performance Comparison: Image Classification Transfer Learning with TensorFlow and the Intel® Transfer Learning Tool\n",
    "\n",
    "This notebook uses the TensorFlow libraries to do transfer learning with an image classification model. The model is exported, evaluated, and used to generate predictions. The same sequence is also done using the Intel Transfer Learning Tool. The Intel Transfer Learning Tool is also used to optimize and quantized the trained model.\n",
    "\n",
    "Graphs are generated to visually compare:\n",
    "* Training metrics (time per epoch, accuracy by epoch, loss by epoch)\n",
    "* Evaluation metrics (time to evaluate the validation dataset, accuracy using the validation data)\n",
    "* Prediction time for a single batch\n",
    "* Latency and throughput for the trained models, quantized model, and the optimized model.\n",
    "\n",
    "The notebook has variables to allow controlling parameters such as the model name, dataset, the number of training epochs, and the batch size(s)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f305fbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' \n",
    "\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as mtick\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import psutil\n",
    "import random\n",
    "import tempfile\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "import warnings\n",
    "\n",
    "from tlt.datasets import dataset_factory\n",
    "from tlt.models import model_factory\n",
    "from tlt.utils.file_utils import download_and_extract_tar_file\n",
    "from tlt.utils.platform_util import CPUInfo, OptimizedPlatformUtil, PlatformUtil\n",
    "from utils import inc_utils\n",
    "\n",
    "# Ignore all warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "# Specify the the default dataset directory\n",
    "dataset_directory = os.environ[\"DATASET_DIR\"] if \"DATASET_DIR\" in os.environ else \\\n",
    "    os.path.join(os.environ[\"HOME\"], \"dataset\")\n",
    "\n",
    "# Specify a directory for output (saved models and checkpoints)\n",
    "output_directory = os.environ[\"OUTPUT_DIR\"] if \"OUTPUT_DIR\" in os.environ else \\\n",
    "    os.path.join(os.environ[\"HOME\"], \"output\")\n",
    "\n",
    "print(\"Output directory:\", output_directory)\n",
    "\n",
    "# TF Hub cache directory\n",
    "os.environ[\"TFHUB_CACHE_DIR\"] = os.path.join(output_directory, \".cache\", \"tfhub_modules\")\n",
    "\n",
    "# Data Frame styles\n",
    "table_styles =[{\n",
    "    'selector': 'caption',\n",
    "    'props': [\n",
    "        ('text-align', 'center'),\n",
    "        ('color', 'black'),\n",
    "        ('font-size', '16px')\n",
    "    ]\n",
    "}]\n",
    "\n",
    "# Colors used in charts\n",
    "orange = '#ff6f00'\n",
    "blue = '#0071c5'\n",
    "dark_blue = '#003c71'\n",
    "yellow = '#f3d54e'\n",
    "\n",
    "# Caption style for DataFrames\n",
    "caption_style = [dict(selector=\"caption\", props=[(\"text-align\", \"center\"), (\"font-size\", \"14pt\"), (\"color\", \"black\")])]\n",
    "\n",
    "# Line styles\n",
    "line_styles = ['solid', 'dotted', 'dashed', 'dashdot']\n",
    "\n",
    "# Marker styles\n",
    "marker_styles = ['o', 'D', 's', 'v']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4270c3e8",
   "metadata": {},
   "source": [
    "## 1. Display Platform Information\n",
    "\n",
    "Use the `CPUInfo` and `PlatformUtil` classes in the get and display information about the platform and TensorFlow version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f3c3ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get and display CPU/platform information\n",
    "cpu_info = CPUInfo()\n",
    "platform_util = PlatformUtil()\n",
    "print(\"{0} CPU Information {0}\".format(\"=\" * 20))\n",
    "print(\"CPU family:\", platform_util.cpu_family)\n",
    "print(\"CPU model:\", platform_util.cpu_model)\n",
    "print(\"CPU type:\", platform_util.cpu_type)\n",
    "print(\"Physical cores per socket:\", cpu_info.cores_per_socket)\n",
    "print(\"Total physical cores:\", cpu_info.cores)\n",
    "cpufreq = psutil.cpu_freq()\n",
    "print(\"Max Frequency:\", cpufreq.max)\n",
    "print(\"Min Frequency:\", cpufreq.min)\n",
    "cpu_socket_count = cpu_info.sockets\n",
    "print(\"Socket Number:\", cpu_socket_count)\n",
    "\n",
    "print(\"\\n{0} Memory Information {0}\".format(\"=\" * 20))\n",
    "svmem = psutil.virtual_memory()\n",
    "print(\"Total: \", int(svmem.total / (1024 ** 3)), \"GB\")\n",
    "\n",
    "# Display TensorFlow version information\n",
    "print(\"\\n{0} TensorFlow Information {0}\".format(\"=\" * 20))\n",
    "print(\"TensorFlow version:\", tf.__version__)\n",
    "print(\"TensorFlow Hub version:\", hub.__version__)\n",
    "major_version = int(tf.__version__.split(\".\")[0])\n",
    "minor_version = int(tf.__version__.split(\".\")[1])\n",
    "if major_version >= 2:\n",
    "    onednn_enabled = 0\n",
    "    if minor_version < 5:\n",
    "        from tensorflow.python import _pywrap_util_port\n",
    "    else:\n",
    "        from tensorflow.python.util import _pywrap_util_port\n",
    "        onednn_enabled = int(os.environ.get('TF_ENABLE_ONEDNN_OPTS', '0'))\n",
    "    on_onednn = _pywrap_util_port.IsMklEnabled() or (onednn_enabled == 1)\n",
    "else:\n",
    "    on_onednn = tf.pywrap_tensorflow.IsMklEnabled()\n",
    "\n",
    "print(\"oneDNN enabled:\", on_onednn)\n",
    "\n",
    "# Don't use the NVidia GPU, if there is one\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d5a1252",
   "metadata": {},
   "source": [
    "## 2. Select a model and define parameters to use during training and evaluation\n",
    "\n",
    "### Select a model\n",
    "\n",
    "See the list of supported image classification models from TensorFlow Hub in the Intel Transfer Learning Tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4d9b894",
   "metadata": {},
   "outputs": [],
   "source": [
    "framework = 'tensorflow'\n",
    "use_case = 'image_classification'\n",
    "model_hub = 'TFHub'\n",
    "supported_models = model_factory.get_supported_models(framework, use_case)\n",
    "supported_models = supported_models[use_case]\n",
    "\n",
    "# Filter to only get relevant models\n",
    "supported_models = { key:value for (key,value) in supported_models.items() if value[framework]['model_hub'] == model_hub}\n",
    "\n",
    "print(\"Supported {} models for {} from {}\".format(framework, use_case, model_hub))\n",
    "print(\"=\" * 70)\n",
    "for model_name in supported_models.keys():\n",
    "    print(model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd14554e",
   "metadata": {},
   "source": [
    "Set the `model_name` to the model that will be used for transfer learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bca71035",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select a model\n",
    "model_name = \"resnet_v1_50\"\n",
    "\n",
    "# Get information about the model (image size and the feature vector handle)\n",
    "# This information will be used during transfer learning using the TensorFlow framework API\n",
    "if model_name in supported_models.keys():\n",
    "    model_info = supported_models[model_name][framework]\n",
    "    image_size = model_info[\"image_size\"]\n",
    "    feature_vector_handle = model_info['feature_vector']\n",
    "    \n",
    "    print(\"Model Name: {}\".format(model_name))\n",
    "    print(\"TF Hub feature vector: {}\".format(feature_vector_handle))\n",
    "    print(\"Image size: {}\".format(image_size))\n",
    "else:\n",
    "    raise ValueError(\"The specified model is unsupported. Please select a model from the list of supported models.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e05fc52e",
   "metadata": {},
   "source": [
    "### Select a dataset\n",
    "\n",
    "By default, the notebook will use the [TensorFlow Flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers), which has flower images that belong to 5 categories.\n",
    "\n",
    "To use your own dataset, set the `dataset_subdir` variable to the dataset path. The dataset directory is expected to have folders of images for each class, where the name of the folder will be used as the class name.\n",
    "\n",
    "```\n",
    "dataset_dir\n",
    "          ├── class_a\n",
    "          ├── class_b\n",
    "          └── class_c\n",
    "```\n",
    "\n",
    "Optionally, the `dataset_subdir` directory can have `train` and `test`/`validation` subdirectories. For example:\n",
    "```\n",
    "dataset_dir\n",
    "          ├── train\n",
    "          |   ├── class_a\n",
    "          |   ├── class_b\n",
    "          |   └── class_c\n",
    "          └── test\n",
    "              ├── class_a\n",
    "              ├── class_b\n",
    "              └── class_c\n",
    "```\n",
    "If the dataset does not have separate folders for train and test/validation, the dataset will be split by percentage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a0d6c0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_subdir = os.path.join(dataset_directory, \"flower_photos\")\n",
    "\n",
    "# Download the flowers dataset, if the folder doesn't exist\n",
    "if not os.path.exists(dataset_subdir):\n",
    "    os.makedirs(dataset_subdir)\n",
    "    dataset_url = \"https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz\"\n",
    "    download_and_extract_tar_file(dataset_url, dataset_directory)\n",
    "    \n",
    "print(\"Dataset path:\", dataset_subdir)\n",
    "\n",
    "print(\"\\nFolders in the dataset directory:\")\n",
    "for d in os.listdir(dataset_subdir):\n",
    "    if os.path.isdir(os.path.join(dataset_subdir, d)):\n",
    "        print(\"-\", d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1aa58df",
   "metadata": {},
   "source": [
    "### Define parameters\n",
    "\n",
    "For consistency between the model training using the TensorFlow framework API and the model training using the Intel Transfer Learning Tool API, the next cell defines parameters that will be used by both methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eee80253",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of training epochs\n",
    "training_epochs = 2\n",
    "\n",
    "# Shuffle the files after each training epoch\n",
    "shuffle_files = True\n",
    "\n",
    "# Define training/validation splits for the dataset \n",
    "# (if the dataset directory does not have subdirectories for train and test/validation)\n",
    "validation_split = 0.25\n",
    "training_split = 1 - validation_split\n",
    "\n",
    "# Set seed for consistency between runs (or None)\n",
    "seed = 10\n",
    "\n",
    "# List of batch size(s) to compare (maximum of 4 batch sizes to try)\n",
    "batch_size_list = [ 256, 512 ]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4847ae1d",
   "metadata": {},
   "source": [
    "Validate parameter values and then print out the parameter values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "051ea2a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not isinstance(training_epochs, int):\n",
    "    raise TypeError(\"The training_epochs parameter should be an integer, but found a {}\".format(type(training_epochs)))\n",
    "\n",
    "if training_epochs < 1:\n",
    "    raise ValueError(\"The training_epochs parameter should not be less than 1.\")\n",
    "    \n",
    "if not isinstance(shuffle_files, bool):\n",
    "    raise TypeError(\"The shuffle_files parameter should be a bool, but found a {}\".format(type(shuffle_files)))\n",
    "\n",
    "if not isinstance(validation_split, float):\n",
    "    raise TypeError(\"The validation_split parameter should be a float, but found a {}\".format(type(validation_split)))\n",
    "\n",
    "if not isinstance(training_split, float):\n",
    "    raise TypeError(\"The training_split parameter should be a float, but found a {}\".format(type(training_split)))\n",
    "\n",
    "if validation_split + training_split > 1:\n",
    "    raise ValueError(\"The sum of validation_split and training_split should not be greater than 1.\")\n",
    "\n",
    "if seed and not isinstance(seed, int):\n",
    "    raise TypeError(\"The seed parameter should be a integer or None, but found a {}\".format(type(seed)))\n",
    "\n",
    "if len(batch_size_list) > 4 or len(batch_size_list) == 0:\n",
    "    raise ValueError(\"The batch_size_list should have at most 4 values, but found {} values ({})\".format(\n",
    "        len(batch_size_list), batch_size_list))\n",
    "    \n",
    "print(\"Number of training epochs:\", training_epochs)\n",
    "print(\"Shuffle files:\", shuffle_files)\n",
    "print(\"Training split: {}%\".format(training_split*100))\n",
    "print(\"Validation split: {}%\".format(validation_split*100))\n",
    "print(\"Seed:\", str(seed))\n",
    "print(\"Batch size list:\", batch_size_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb2d786a",
   "metadata": {},
   "source": [
    "Define a callback method that track the time that it took to run training epochs, evaluation, and batch predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "713f648f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Callback to track the training time for each epoch, evaluation time, or prediction time\n",
    "class TimerCallback(tf.keras.callbacks.Callback):\n",
    "    def __init__(self):\n",
    "        self.epoch_times = []\n",
    "        self.eval_times = []\n",
    "        self.predict_times = []\n",
    "    def on_epoch_begin(self, batch, logs={}):\n",
    "        self.tf_timestamp = tf.timestamp()\n",
    "    def on_epoch_end(self,epoch,logs = {}):\n",
    "        self.epoch_times.append((tf.timestamp() - self.tf_timestamp).numpy())\n",
    "    def on_test_begin(self, batch, logs={}):\n",
    "        self.tf_timestamp = tf.timestamp()\n",
    "    def on_test_end(self,epoch,logs = {}):\n",
    "        self.eval_times.append((tf.timestamp() - self.tf_timestamp).numpy())\n",
    "    def on_predict_begin(self, batch, logs={}):\n",
    "        self.tf_timestamp = tf.timestamp()\n",
    "    def on_predict_end(self,epoch,logs = {}):\n",
    "        self.predict_times.append((tf.timestamp() - self.tf_timestamp).numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd401a7",
   "metadata": {},
   "source": [
    "## 3. Compare the training time for transfer learning\n",
    "\n",
    "In this section, we will compare the time it takes to retrain the image classification model using the dataset that was selected in the previous section.\n",
    "\n",
    "The training will be done in two different ways to compare:\n",
    "* Transfer learning using the TensorFlow framework and TF Hub libraries\n",
    "* Transfer learning using the Intel Transfer Learning Tool API\n",
    "\n",
    "### Transfer learning using the TensorFlow framework and TF Hub libraries\n",
    "\n",
    "This section goes through using the TensorFlow framework and TF Hub libraries to retrain the model using the selected dataset.\n",
    "\n",
    "First, the dataset is loaded in, which allows us to determine the number of classes in the dataset. The original ImageNet dataset that the image classification model was trained on has 1000 classes. To do transfer learning using the new dataset, we will get the feature vector from TF Hub and then add on a classification layer that matches the number of classes in the new dataset.\n",
    "\n",
    "If multiple batch sizes were set in the `batch_size_list`, training will be run for each batch size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08cfd3ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set seed\n",
    "if seed:\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    tf.random.set_seed(seed)\n",
    "\n",
    "# Lists to track callbacks, datasets, models, and saved model directory for each batch size experiment\n",
    "tf_time_callback_list = []\n",
    "tf_dataset_list = []\n",
    "tf_model_list = []\n",
    "tf_export_dir_list = []\n",
    "tf_history_list = []\n",
    "\n",
    "# Check if the dataset directory has subdirectories for train/validation/test splits\n",
    "val_dataset_dir = None\n",
    "train_dataset_dir = dataset_subdir\n",
    "if os.path.exists(os.path.join(dataset_subdir, 'train')):\n",
    "    train_dataset_dir = os.path.join(dataset_subdir, 'train')\n",
    "    val_dataset_dir = os.path.join(dataset_subdir, 'validation')\n",
    "    \n",
    "    if not os.path.exists(val_dataset_dir):\n",
    "        if os.path.exists(os.path.join(dataset_subdir, 'test')):\n",
    "            val_dataset_dir = os.path.join(dataset_subdir, 'test')\n",
    "        else:\n",
    "            raise ValueError('The dataset directory ({}) has a \"train\" directory, but no \"validation\" or \"test\" directory.')\n",
    "\n",
    "    print(\"Using training data from {}\".format(train_dataset_dir))\n",
    "    print(\"Using validation data from {}\".format(val_dataset_dir))\n",
    "    \n",
    "# Load the dataset\n",
    "tf_dataset = tf.keras.utils.image_dataset_from_directory(train_dataset_dir, batch_size=None, seed=seed)\n",
    "class_names = tf_dataset.class_names\n",
    "\n",
    "if shuffle_files:\n",
    "    tf_dataset = tf_dataset.shuffle(tf_dataset.cardinality(), reshuffle_each_iteration=False, seed=seed)\n",
    "\n",
    "if val_dataset_dir:\n",
    "    # Load the validation/test sub directory\n",
    "    train_ds = tf_dataset\n",
    "    val_ds = tf.keras.utils.image_dataset_from_directory(val_dataset_dir, batch_size=None, seed=seed) \n",
    "    if shuffle_files:\n",
    "        val_ds = val_ds.shuffle(val_ds.cardinality(), reshuffle_each_iteration=False, seed=seed)\n",
    "    train_size = len(train_ds)\n",
    "    val_size = len(val_ds)\n",
    "else:\n",
    "    # Split the data into train/validation subsets (Note that image_dataset_from_directory can also do splitting but\n",
    "    # we are doing it this way to match what the Intel Transfer Learning Tool does to ensure the same sized splits)\n",
    "    train_size = int(training_split * len(tf_dataset))\n",
    "    val_size = int(validation_split * len(tf_dataset))\n",
    "    train_ds = tf_dataset.take(train_size)\n",
    "    val_ds = tf_dataset.skip(train_size).take(val_size)\n",
    "\n",
    "print(\"Training dataset size:\", train_size)\n",
    "print(\"Validation dataset size:\", val_size)\n",
    "    \n",
    "# Preprocess the dataset\n",
    "normalization_layer = tf.keras.layers.Rescaling(1. / 255)\n",
    "\n",
    "def preprocess_image(image, label):\n",
    "    image = tf.image.resize_with_pad(image, image_size, image_size)\n",
    "    image = normalization_layer(image)\n",
    "    return (image, label)\n",
    "\n",
    "train_ds = train_ds.map(preprocess_image)\n",
    "val_ds = val_ds.map(preprocess_image)\n",
    "\n",
    "for batch_size in batch_size_list:\n",
    "    print('-' * 40)\n",
    "    print('Training using batch size: {}'.format(batch_size))\n",
    "    print('-' * 40)\n",
    "    \n",
    "    # Batch the dataset \n",
    "    batched_train_ds = train_ds.batch(batch_size)\n",
    "    batched_val_ds = val_ds.batch(batch_size)\n",
    "    \n",
    "    # Get the feature extractor layer from TF Hub\n",
    "    feature_extractor_layer = hub.KerasLayer(feature_vector_handle,\n",
    "                                             input_shape=(image_size, image_size, 3),\n",
    "                                             trainable=False)\n",
    "\n",
    "    # Add the dense layer sized according to the number of classes in our dataset\n",
    "    tf_model = tf.keras.Sequential([\n",
    "      feature_extractor_layer,\n",
    "      tf.keras.layers.Dense(len(class_names))\n",
    "    ])\n",
    "\n",
    "    # Configure the model optimizer and loss function\n",
    "    tf_model.compile(\n",
    "      optimizer=tf.keras.optimizers.Adam(),\n",
    "      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "      metrics=['acc'])\n",
    "    \n",
    "    tf_model.summary()\n",
    "\n",
    "    # Define the callback for tracking the time it takes to train each epoch\n",
    "    tf_time_callback = TimerCallback()\n",
    "\n",
    "    # Train the model\n",
    "    tf_history_list.append(tf_model.fit(batched_train_ds, epochs=training_epochs, shuffle=shuffle_files,\n",
    "                                        callbacks=[tf_time_callback]))\n",
    "    \n",
    "    # Export the trained model\n",
    "    tf_export_dir = os.path.join(output_directory, \"tf_saved_models\", model_name)\n",
    "    if not os.path.exists(tf_export_dir):\n",
    "        os.makedirs(tf_export_dir)\n",
    "    tf_export_dir = tempfile.mkdtemp(prefix=tf_export_dir + '/')\n",
    "    print(\"Save model to:\", tf_export_dir)\n",
    "    tf_model.save(tf_export_dir)\n",
    "    \n",
    "    # Append to lists for each batch size\n",
    "    tf_time_callback_list.append(tf_time_callback)\n",
    "    tf_dataset_list.append((batched_train_ds, batched_val_ds))\n",
    "    tf_model_list.append(tf_model)\n",
    "    tf_export_dir_list.append(tf_export_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "160cf79a",
   "metadata": {},
   "source": [
    "### Transfer learning using the Intel Transfer Learning Tool API\n",
    "\n",
    "This section uses the Intel Transfer Learning Tool API to retrain the model using the selected dataset. This API simplifies the transfer learning process, so there are less lines of code compared to directly using the TensorFlow and TensorFlow Hub libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fd15b3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the OptimizedPlatform Util class from the Intel Transfer Learning Tool API to set recommended settings\n",
    "optimized_platform_util = OptimizedPlatformUtil(omp_num_threads=cpu_info.cores_per_socket,\n",
    "                                               kmp_blocktime=0,\n",
    "                                               kmp_affinity='granularity=fine,compact,1,0',\n",
    "                                               tf_num_intraop_threads=cpu_info.cores_per_socket,\n",
    "                                               tf_num_interop_threads=cpu_info.sockets,\n",
    "                                               force_reset_env_vars=True)\n",
    "\n",
    "for k, v in optimized_platform_util.env_vars_dict.items():\n",
    "    if v is not None:\n",
    "        print(\"{}: {}\".format(k, v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fac3d6df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lists to track callbacks, datasets, models, and saved model directory for each batch size experiment\n",
    "tlt_time_callback_list = []\n",
    "tlt_dataset_list = []\n",
    "tlt_model_list = []\n",
    "tlt_export_dir_list = []\n",
    "tlt_history_list = []\n",
    "    \n",
    "for batch_size in batch_size_list:\n",
    "    print('-' * 40)\n",
    "    print('Training using batch size: {}'.format(batch_size))\n",
    "    print('-' * 40)\n",
    "    \n",
    "    # Use the model factory to get the model\n",
    "    tlt_model = model_factory.get_model(model_name=model_name, framework=framework)\n",
    "    \n",
    "    # Load, split, and preprocess the dataset\n",
    "    tlt_dataset = dataset_factory.load_dataset(dataset_dir=dataset_subdir, use_case=use_case, framework=framework)\n",
    "    \n",
    "    if not tlt_dataset.train_subset:\n",
    "        tlt_dataset.shuffle_split(train_pct=training_split, val_pct=validation_split, seed=seed, shuffle_files=shuffle_files)\n",
    "    \n",
    "    tlt_dataset.preprocess(tlt_model.image_size, batch_size=batch_size)\n",
    "    \n",
    "    # Define the callback for tracking the time it takes to train each epoch\n",
    "    tlt_time_callback = TimerCallback()\n",
    "\n",
    "    # Train the model\n",
    "    tlt_history_list.append(tlt_model.train(tlt_dataset, output_dir=output_directory, epochs=training_epochs,\n",
    "                                            shuffle_files=shuffle_files, do_eval=False, callbacks=tlt_time_callback,\n",
    "                                            seed=seed))\n",
    "\n",
    "    # Export the trained model\n",
    "    tlt_export_dir = os.path.join(output_directory, \"tlt_saved_models\")\n",
    "    tlt_export_dir = tlt_model.export(tlt_export_dir)\n",
    "    \n",
    "    # Append to lists for each batch size\n",
    "    tlt_time_callback_list.append(tlt_time_callback)\n",
    "    tlt_dataset_list.append(tlt_dataset)\n",
    "    tlt_model_list.append(tlt_model)\n",
    "    tlt_export_dir_list.append(tlt_export_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5e52a50",
   "metadata": {},
   "source": [
    "### Optimize the model using the Intel Transfer Learning Tool API\n",
    "\n",
    "After training, the Intel Transfer Learning Tool can optimize the model to improve inference performance. This is done using the [Intel® Neural Compressor](https://github.com/intel/neural-compressor) quantizing the model or optimizing the full precision model. \n",
    "\n",
    "First, we setup a configuration file that with parameters that will be used by the Intel Neural Compressor for quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df079cb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "tlt_quantization_dir_list = []\n",
    "tlt_optimized_dir_list = []\n",
    "inc_config_list = []\n",
    "\n",
    "# Create a tuning workspace directory for INC\n",
    "nc_workspace = os.path.join(output_directory, 'nc_workspace')\n",
    "\n",
    "# Relative accuracy loss (1%)\n",
    "relative_accuracy_criterion = 0.01\n",
    "\n",
    "# Define the exit policy timeout (in seconds) and max number of trials. The tuning processing finishes when\n",
    "# the timeout or max trials is reached. A tuning timeout of 0 means that the tuning phase stops when the\n",
    "# accuracy criterion is met.\n",
    "timeout = 0\n",
    "max_trials=15\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    # Create an output directories for the quantized and optimized models\n",
    "    tlt_quantization_dir = os.path.join(output_directory, 'tlt_quantized_models', model_name, os.path.basename(tlt_export_dir_list[i]))\n",
    "    tlt_optimized_dir = os.path.join(output_directory, 'tlt_optimized_models', model_name, os.path.basename(tlt_export_dir_list[i]))\n",
    "\n",
    "    # Create an Intel Neural Compressor config based on the inputs that we are using\n",
    "    inc_config_list.append(tlt_model.get_inc_config(accuracy_criterion_relative=relative_accuracy_criterion,\n",
    "                                                    exit_policy_timeout=timeout, exit_policy_max_trials=max_trials))\n",
    "    \n",
    "    # Append to lists for each batch size\n",
    "    tlt_quantization_dir_list.append(tlt_quantization_dir)\n",
    "    tlt_optimized_dir_list.append(tlt_optimized_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "534de66e",
   "metadata": {},
   "source": [
    "Next, we quantize the model using the config file that was just generated. Quantization aims to improve inference\n",
    "performance by reducing the number of bits required, by maintaining close the the same amount of accuracy as the full precision model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103ffd57",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    # Quantize the model\n",
    "    tlt_model.quantize(tlt_quantization_dir_list[i], tlt_dataset_list[i], config=inc_config_list[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfeea7eb",
   "metadata": {},
   "source": [
    "Another option to improve inference performance is using graph optimization through the Intel Neural Compressor which:\n",
    "* Converts variables to constants\n",
    "* Removes training-only operations like checkpoint saving\n",
    "* Strips out parts of the graph that are never reached\n",
    "* Removes debug operations like CheckNumerics\n",
    "* Folds batch normalization ops into the pre-calculated weights\n",
    "* Fuses common operations into unified versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a2f7a8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    # Optimize the full precision model\n",
    "    tlt_model.optimize_graph(tlt_optimized_dir_list[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "304258ec",
   "metadata": {},
   "source": [
    "### Compare training times\n",
    "\n",
    "The table below compares the time it took to train each epoch (in seconds) using the TensorFlow framework libraries directly versus the Intel Transfer Learning Tool API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16299291",
   "metadata": {},
   "outputs": [],
   "source": [
    "display_df = []\n",
    "plt.figure(figsize=(10,6))\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    # Sanity check that both datasets had the same number of batches\n",
    "    if len(tf_dataset_list[i][0]) != len(tlt_dataset_list[i].train_subset):\n",
    "        print(\"WARNING: For batch size {}, the TF training dataset had {} batches and the TLT training dataset had \"\n",
    "              \"{} batches. These values should have been the same.\".format(batch_size, len(tf_dataset_list[i][0]), len(tlt_dataset_list[i].train_subset)))\n",
    "    \n",
    "    # Calculate images/second\n",
    "    tf_images_per_second = [train_size / t for t in tf_time_callback_list[i].epoch_times]\n",
    "    tlt_images_per_second = [train_size / t for t in tlt_time_callback_list[i].epoch_times]\n",
    "    performance_delta = [\"{0:.2f}%\".format((tlt-tf)/tf * 100) for tf, tlt in zip(tf_images_per_second, tlt_images_per_second)]\n",
    "\n",
    "    # Graph the results\n",
    "    epoch_list = [str(i) for i in range(1, training_epochs + 1)]\n",
    "    tf_train_time = tf_time_callback_list[i].epoch_times\n",
    "    tlt_train_time = tlt_time_callback_list[i].epoch_times\n",
    "\n",
    "    plt.plot(epoch_list, tf_train_time, label=\"Using TF libraries with batch size {}\".format(batch_size),\n",
    "             linestyle=line_styles[i], marker=marker_styles[i], color=orange)\n",
    "    plt.plot(epoch_list, tlt_train_time, label=\"Using TLT with batch size {}\".format(batch_size), \n",
    "             linestyle=line_styles[i], marker=marker_styles[i],color=blue)\n",
    "    \n",
    "    # Create a DataFrame to display the results in a table\n",
    "    df = pd.DataFrame({\n",
    "        'TF epoch time<br>(seconds)': tf_time_callback_list[i].epoch_times,\n",
    "        'TLT epoch time<br>(seconds)': tlt_time_callback_list[i].epoch_times,\n",
    "        'TF throughput<br>(images/sec)': tf_images_per_second,\n",
    "        'TLT throughput<br>(images/sec)': tlt_images_per_second,\n",
    "        'Performance<br>Boost': performance_delta\n",
    "    })\n",
    "    df.index += 1 \n",
    "    df = df.style.set_table_styles(table_styles).set_caption(\"Epoch training times with batch size {}\".format(batch_size))\n",
    "    display_df.append(df)\n",
    "\n",
    "plt.title(\"Training time per epoch\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Seconds\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# Display tables with epoch training time for each batch size\n",
    "for df in display_df:\n",
    "    display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "388a7675",
   "metadata": {},
   "source": [
    "Next, visualize the accuracy and loss metrics collected during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59598b38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Graph the training accuracy by epoch for each batch size\n",
    "plt.figure(figsize=(10,6))\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    tf_acc_time = [i * 100 for i in tf_history_list[i].history['acc']]\n",
    "    tlt_acc_time = [i * 100 for i in tlt_history_list[i]['acc']]\n",
    "\n",
    "    plt.plot(epoch_list, tf_acc_time, label = \"Using TF libraries (batch size = {})\".format(batch_size), linestyle=line_styles[i], marker=marker_styles[i], color=orange)\n",
    "    plt.plot(epoch_list, tlt_acc_time, label = \"Using TLT (batch size = {})\".format(batch_size), linestyle=line_styles[i], marker=marker_styles[i], color=blue)\n",
    "\n",
    "plt.title(\"Training Accuracy by Epoch\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Accuracy (%)\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# Graph the training loss by epoch for each batch size\n",
    "plt.figure(figsize=(10,6))\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    tf_loss_time = tf_history_list[i].history['loss']\n",
    "    tlt_loss_time = tlt_history_list[i]['loss']\n",
    "\n",
    "    plt.plot(epoch_list, tf_loss_time, label = \"Using TF libraries (batch size = {})\".format(batch_size), linestyle=line_styles[i], marker=marker_styles[i], color=orange)\n",
    "    plt.plot(epoch_list, tlt_loss_time, label = \"Using TLT (batch size = {})\".format(batch_size), linestyle=line_styles[i], marker=marker_styles[i], color=blue)\n",
    "\n",
    "plt.title(\"Training Loss by Epoch\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e58715d4",
   "metadata": {},
   "source": [
    "## 4. Evaluate and predict\n",
    "\n",
    "This section calls evaluation and prediction methods for the models trained using the TensorFlow libraries and the Intel Transfer Learning Tool.\n",
    "\n",
    "### Evaluate the models using the validation data\n",
    "\n",
    "First, evaluate the models trained using the TensorFlow libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba3aadb",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_eval_callback_list = []\n",
    "tf_eval_metrics_list = []\n",
    "\n",
    "# Evaluate using the TensorFlow framework model for each batch size\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    print('-' * 40)\n",
    "    print('Evaluate using batch size: {}'.format(batch_size))\n",
    "    print('-' * 40)\n",
    "    \n",
    "    tf_eval_callback = TimerCallback()\n",
    "    \n",
    "    # Use the test split of the dataset to evaluate the model\n",
    "    tf_eval_metrics_list.append(tf_model_list[i].evaluate(tf_dataset_list[i][1], callbacks=tf_eval_callback))\n",
    "    tf_eval_callback_list.append(tf_eval_callback)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de6dbc91",
   "metadata": {},
   "source": [
    "Next, evaluate the models trained using the Intel Transfer Learning Tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed9e1a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "tlt_eval_callback_list = []\n",
    "tlt_eval_metrics_list = []\n",
    "\n",
    "# Evaluate using the Intel Transfer Learning Tool model for each batch size\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    print('-' * 40)\n",
    "    print('Evaluate using batch size: {}'.format(batch_size))\n",
    "    print('-' * 40)\n",
    "    \n",
    "    use_test_set = tlt_dataset_list[i].validation_subset is None and tlt_dataset_list[i].test_subset is not None\n",
    "    \n",
    "    tlt_eval_callback = TimerCallback()\n",
    "    tlt_eval_metrics_list.append(tlt_model_list[i].evaluate(tlt_dataset_list[i], callbacks=tlt_eval_callback, use_test_set=use_test_set))\n",
    "    tlt_eval_callback_list.append(tlt_eval_callback)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "231f3731",
   "metadata": {},
   "source": [
    "After all the models have been evaluated, visualize the results using charts that the display the time that it took to evaluate each model and the accuracy that was found when using the validation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c31fee1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bar chart group labels\n",
    "groups = [\"batch size = {}\".format(bs) for bs in batch_size_list]\n",
    "\n",
    "# Create grouped bar chart for evaluation time\n",
    "decimals = 3 # number of decimals to use for rounding\n",
    "tf_eval_times = [round(callback.eval_times[0], decimals) for callback in tf_eval_callback_list]\n",
    "tlt_eval_times = [round(callback.eval_times[0], decimals) for callback in tlt_eval_callback_list]\n",
    "\n",
    "x = np.arange(len(groups))\n",
    "width = 0.24  # the width of the bars\n",
    "multiplier = 0\n",
    "\n",
    "# Setup bars for evaluation times\n",
    "fig, (ax1, ax2) = plt.subplots(2)\n",
    "fig.set_figheight(10)\n",
    "fig.set_figwidth(10)\n",
    "rects_tf = ax1.bar(x - width/2, tf_eval_times, width, label='TF eval', color=orange)\n",
    "rects_tlt = ax1.bar(x + width/2, tlt_eval_times, width, label='TLT eval', color=blue)\n",
    "ax1.bar_label(rects_tf, padding=3)\n",
    "ax1.bar_label(rects_tlt, padding=3)\n",
    "\n",
    "# Add labels, title, and legend\n",
    "ax1.set_ylabel('Seconds')\n",
    "ax1.set_title('Evaluation time')\n",
    "ax1.set_xticks(x, groups)\n",
    "ax1.set_ymargin(0.2) \n",
    "ax1.legend(ncols=2)\n",
    "#plt.show()\n",
    "\n",
    "# Evaluation accuracy comparison\n",
    "decimals = 2\n",
    "tf_acc_index = tf_model_list[0].metrics_names.index('acc')\n",
    "tlt_acc_index = tlt_model_list[0]._model.metrics_names.index('acc')\n",
    "tf_eval_accuracy = [round(x[tf_acc_index] * 100, decimals) for x in tf_eval_metrics_list]\n",
    "tlt_eval_accuracy = [round(x[tlt_acc_index] * 100, decimals) for x in tlt_eval_metrics_list]\n",
    "\n",
    "# Setup bars for evaluation accuracy\n",
    "rects_tf = ax2.bar(x - width/2, tf_eval_accuracy, width, label='TF accuracy', color=orange)\n",
    "rects_tlt = ax2.bar(x + width/2, tlt_eval_accuracy, width, label='TLT accuracy', color=blue)\n",
    "ax2.bar_label(rects_tf, padding=3)\n",
    "ax2.bar_label(rects_tlt, padding=3)\n",
    "\n",
    "# Add labels, title, and legend\n",
    "ax2.set_ylabel('Accuracy (%)')\n",
    "ax2.yaxis.set_major_formatter(mtick.PercentFormatter())\n",
    "ax2.set_title('Evaluation accuracy using the validation data')\n",
    "ax2.set_xticks(x, groups)\n",
    "ax2.set_ymargin(0.2) \n",
    "ax2.legend(ncols=2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d759b5d",
   "metadata": {},
   "source": [
    "### Predict using a batch of images\n",
    "\n",
    "Use the TensorFlow libaries to get a batch of images and predict using the trained models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38132396",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_predict_callback_list = []\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    print('-' * 50)\n",
    "    print('Predict on a single batch (batch size = {})'.format(batch_size))\n",
    "    print('-' * 50)\n",
    "    \n",
    "    tf_predict_time = TimerCallback()\n",
    "    dataset_batch = next(iter(tf_dataset_list[i][0]))\n",
    "    tf_batch, _ = dataset_batch\n",
    "    batch_predictions = tf_model_list[i].predict(tf_batch, callbacks=tf_predict_time)\n",
    "    tf_predict_callback_list.append(tf_predict_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d3beb10",
   "metadata": {},
   "source": [
    "Similarly, use the Intel Transfer Learning Tool API to get a batch of images and predict using the trained models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66be9e10",
   "metadata": {},
   "outputs": [],
   "source": [
    "tlt_predict_callback_list = []\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    print('-' * 50)\n",
    "    print('Predict on a single batch (batch size = {})'.format(batch_size))\n",
    "    print('-' * 50)\n",
    "    \n",
    "    tlt_predict_time = TimerCallback()\n",
    "\n",
    "    tlt_batch, _ = tlt_dataset_list[i].get_batch(subset='train')\n",
    "    predictions = tlt_model_list[i].predict(tlt_batch, callbacks=tlt_predict_time)\n",
    "    tlt_predict_callback_list.append(tlt_predict_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a594631f",
   "metadata": {},
   "source": [
    "Visualize the time that it took to get predictions for a batch of images for each model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "309f636a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create grouped bar chart for prediction time\n",
    "decimals = 3   # number of decimals to use for rounding\n",
    "tf_predict_times = [round(callback.predict_times[0], decimals) for callback in tf_predict_callback_list]\n",
    "tlt_predict_times = [round(callback.predict_times[0], decimals) for callback in tlt_predict_callback_list]\n",
    "\n",
    "# Setup bars for evaluation times\n",
    "fig, ax = plt.subplots()\n",
    "fig.set_figheight(6)\n",
    "fig.set_figwidth(10)\n",
    "rects_tf = ax.bar(x - width/2, tf_predict_times, width, label='TF predict', color=orange)\n",
    "rects_tlt = ax.bar(x + width/2, tlt_predict_times, width, label='TLT predict', color=blue)\n",
    "ax.bar_label(rects_tf, padding=3)\n",
    "ax.bar_label(rects_tlt, padding=3)\n",
    "\n",
    "# Add labels, title, and legend\n",
    "ax.set_ylabel('Seconds')\n",
    "ax.set_title('Prediction time for a single batch')\n",
    "ax.set_xticks(x, groups)\n",
    "ax.set_ymargin(0.2) \n",
    "ax.legend(ncols=2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca0771ec",
   "metadata": {},
   "source": [
    "### Check performance using the Intel® Neural Compressor\n",
    "\n",
    "Use the [Intel Neural Compressor](https://github.com/intel/neural-compressor/tree/master) to determine the performance of the exported models. \n",
    "\n",
    "We will compare:\n",
    "* The original model that was trained using the TensorFlow and TF Hub libaries\n",
    "* The model trained using the Intel Transfer Learning Tool\n",
    "* The model trained and quantized using the Intel Transfer Learning Tool\n",
    "* The model trained and optimized using the Intel Transfer Learning Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41b6baee",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset_dir = dataset_subdir\n",
    "\n",
    "if os.path.exists(os.path.join(dataset_subdir, 'validation')):\n",
    "    test_dataset_dir = os.path.join(dataset_subdir, 'validation')\n",
    "elif os.path.exists(os.path.join(dataset_subdir, 'test')):\n",
    "    test_dataset_dir = os.path.join(dataset_subdir, 'test')\n",
    "    \n",
    "print(\"Test dataset directory:\", test_dataset_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e89ccbfb",
   "metadata": {},
   "source": [
    "Use the Intel Neural Compressor to get the performance of the model trained using the TensorFlow libraries.\n",
    "\n",
    "Note that you may see a `zmq.error.ZMQError: Address already in use` error in the output, which is a known issuen when running the Intel Neural Compressor from Jupyter notebooks. If this happens, rerun the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80be966e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf_latency_list = []\n",
    "tf_throughput_list = []\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    print('-' * 90)\n",
    "    print('Check performance for TF model (batch size = {})'.format(batch_size))\n",
    "    print('Saved model directory: {}'.format(tf_export_dir_list[i]))\n",
    "    print('-' * 90)\n",
    "    \n",
    "    results = inc_utils.performance(tf_export_dir_list[i], batch_size, image_size, test_dataset_dir, framework)\n",
    "    tf_latency, tf_throughput = inc_utils.calculate_latency_and_throughput(results)\n",
    "    \n",
    "    tf_latency_list.append(tf_latency)\n",
    "    tf_throughput_list.append(tf_throughput)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ee14994",
   "metadata": {},
   "source": [
    "Next, get the performance of the the model that was trained and exported by the Intel Transfer Learning Toolkit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb33d150",
   "metadata": {},
   "outputs": [],
   "source": [
    "tlt_latency_list = []\n",
    "tlt_throughput_list = []\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    print('-' * 90)\n",
    "    print('Check performance for TLT model (batch size = {})'.format(batch_size))\n",
    "    print('Saved model directory: {}'.format(tlt_export_dir_list[i]))\n",
    "    print('-' * 90)\n",
    "    \n",
    "    tlt_results = inc_utils.performance(tlt_export_dir_list[i], batch_size, image_size, test_dataset_dir, framework)\n",
    "    tlt_latency, tlt_throughput = inc_utils.calculate_latency_and_throughput(tlt_results)\n",
    "    \n",
    "    tlt_latency_list.append(tlt_latency)\n",
    "    tlt_throughput_list.append(tlt_throughput)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aaf38ec",
   "metadata": {},
   "source": [
    "Get the performance of the model that was quantized using the Intel Transfer Learning tool with the Intel Neural Compressor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7322f6a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "quantized_latency_list = []\n",
    "quantized_throughput_list = []\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    try:\n",
    "        tlt_quantized_latency = 0\n",
    "        tlt_quantized_throughput = 0\n",
    "        \n",
    "        print('-' * 90)\n",
    "        print('Check performance for TLT quantized model (batch size = {})'.format(batch_size))\n",
    "        print('Saved model directory: {}'.format(tlt_quantization_dir_list[i]))\n",
    "        print('-' * 90)\n",
    "        \n",
    "        if not os.path.exists(os.path.join(tlt_quantization_dir_list[i], 'saved_model.pb')):\n",
    "            raise FileNotFoundError(\"The quantized model was not found at: {}\\nQuantization may have failed for this model/batch size.\".format(tlt_quantization_dir_list[i],))\n",
    "    \n",
    "        tlt_quantized_results = inc_utils.performance(tlt_quantization_dir_list[i], batch_size, image_size, test_dataset_dir, framework)\n",
    "        tlt_quantized_latency, tlt_quantized_throughput = inc_utils.calculate_latency_and_throughput(tlt_quantized_results)\n",
    "    except Exception as e:\n",
    "        print(\"Error when trying to check the performance for the quantized model with batch size {}\".format(batch_size))\n",
    "        print(e)\n",
    "    finally:\n",
    "        quantized_latency_list.append(tlt_quantized_latency)\n",
    "        quantized_throughput_list.append(tlt_quantized_throughput)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e262895d",
   "metadata": {},
   "source": [
    "Finally, get the performance of the model that was optimized using the Intel Transfer Learning tool with the Intel Neural Compressor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c42c815",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "optimized_latency_list = []\n",
    "optimized_throughput_list = []\n",
    "\n",
    "for i, batch_size in enumerate(batch_size_list):\n",
    "    try:\n",
    "        tlt_optimized_latency = 0\n",
    "        tlt_optimized_throughput = 0\n",
    "        \n",
    "        print('-' * 90)\n",
    "        print('Check performance for TLT optimized model (batch size = {})'.format(batch_size))\n",
    "        print('Saved model directory: {}'.format(tlt_optimized_dir_list[i]))\n",
    "        print('-' * 90)\n",
    "        \n",
    "        if not os.path.exists(os.path.join(tlt_optimized_dir_list[i], 'saved_model.pb')):\n",
    "            raise FileNotFoundError(\"The optimized model was not found at: {}\\nOptimization may have failed for this model/batch size.\".format(tlt_optimized_dir_list[i],))\n",
    "\n",
    "        tlt_optimized_results = inc_utils.performance(tlt_optimized_dir_list[i], batch_size, image_size, test_dataset_dir, framework)\n",
    "        \n",
    "        tlt_optimized_latency, tlt_optimized_throughput = inc_utils.calculate_latency_and_throughput(tlt_optimized_results)\n",
    "    except Exception as e:\n",
    "        print(\"Error when trying to check the performance for the optimized model with batch size {}\".format(batch_size))\n",
    "        print(e)\n",
    "    finally:\n",
    "        optimized_latency_list.append(tlt_optimized_latency)\n",
    "        optimized_throughput_list.append(tlt_optimized_throughput)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44b337ab",
   "metadata": {},
   "source": [
    "Visualize the latency and throughput results for all of the models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b5ca789",
   "metadata": {},
   "outputs": [],
   "source": [
    "width = 0.18   # the width of the bars\n",
    "\n",
    "# Round the latency values\n",
    "decimals = 2   # number of decimals to use for rounding\n",
    "tf_latency_list = [0 if math.isnan(x) else round(x, decimals) for x in tf_latency_list]\n",
    "tlt_latency_list = [0 if math.isnan(x) else round(x, decimals) for x in tlt_latency_list]\n",
    "quantized_latency_list = [0 if math.isnan(x) else round(x, decimals) for x in quantized_latency_list]\n",
    "optimized_latency_list = [0 if math.isnan(x) else round(x, decimals) for x in optimized_latency_list]\n",
    "\n",
    "# Setup the grouped bar chart for latency\n",
    "fig, ax = plt.subplots()\n",
    "fig.set_figheight(6)\n",
    "fig.set_figwidth(10)\n",
    "rects_tf = ax.bar(x, tf_latency_list, width, label='TF latency', color=orange)\n",
    "rects_tlt = ax.bar(x + width, tlt_latency_list, width, label='TLT latency', color=blue)\n",
    "rects_quant = ax.bar(x + width * 2, quantized_latency_list, width, label='TLT quantized latency', color=yellow)\n",
    "rects_opt = ax.bar(x + width * 3, optimized_latency_list, width, label='TLT optimized latency', color=dark_blue)\n",
    "ax.bar_label(rects_tf, padding=3)\n",
    "ax.bar_label(rects_tlt, padding=3)\n",
    "ax.bar_label(rects_quant, padding=3)\n",
    "ax.bar_label(rects_opt, padding=3)\n",
    "\n",
    "# Add labels, title, and legend\n",
    "ax.set_ylabel('Milliseconds')\n",
    "ax.set_title('Latency')\n",
    "ax.set_xticks(x + width*1.5, groups)\n",
    "ax.set_ymargin(0.2) \n",
    "ax.legend(ncols=2)\n",
    "plt.show()\n",
    "\n",
    "# Round the throughput values\n",
    "decimals = 0   # number of decimals to use for rounding\n",
    "tf_throughput_list = [round(x, decimals) for x in tf_throughput_list]\n",
    "tlt_throughput_list = [round(x, decimals) for x in tlt_throughput_list]\n",
    "quantized_throughput_list = [round(x, decimals) for x in quantized_throughput_list]\n",
    "optimized_throughput_list = [round(x, decimals) for x in optimized_throughput_list]\n",
    "\n",
    "# Setup the grouped bar chart for throughput\n",
    "fig, ax = plt.subplots()\n",
    "fig.set_figheight(6)\n",
    "fig.set_figwidth(10)\n",
    "rects_tf = ax.bar(x, tf_throughput_list, width, label='TF throughput', color=orange)\n",
    "rects_tlt = ax.bar(x + width, tlt_throughput_list, width, label='TLT throughput', color=blue)\n",
    "rects_quant = ax.bar(x + width * 2, quantized_throughput_list, width, label='TLT quantized throughput', color=yellow)\n",
    "rects_opt = ax.bar(x + width * 3, optimized_throughput_list, width, label='TLT optimized throughput', color=dark_blue)\n",
    "ax.bar_label(rects_tf, padding=3)\n",
    "ax.bar_label(rects_tlt, padding=3)\n",
    "ax.bar_label(rects_quant, padding=3)\n",
    "ax.bar_label(rects_opt, padding=3)\n",
    "\n",
    "# Add labels, title, and legend\n",
    "ax.set_ylabel('images/second')\n",
    "ax.set_title('Throughput')\n",
    "ax.set_xticks(x + width*1.5, groups)\n",
    "ax.set_ymargin(0.2) \n",
    "ax.legend(ncols=2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa9644c8",
   "metadata": {},
   "source": [
    "The experiments done in this notebook allowed us to compare the training time and inference/evaluation metrics when using the TensorFlow libraries and the Intel Transfer Learning tool. We can also see how batch size affects performance. More experiments can be done by rerunning this notebook with a different model, different dataset, and/or different training parameters.\n",
    "\n",
    "Other related notebooks:\n",
    "* [Transfer Learning for Image Classification using TensorFlow and the Intel® Transfer Learning Tool API](../image_classification/tlt_api_tf_image_classification/TLT_TF_Image_Classification_Transfer_Learning.ipynb)\n",
    "* [Transfer Learning for Image Classification using PyTorch and the Intel® Transfer Learning Tool API](../image_classification/tlt_api_pyt_image_classification/TLT_PyTorch_Image_Classification_Transfer_Learning.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}