classification_tf2/byom_converters/byom_converter_mobilenetv2.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "402b7a52",
   "metadata": {},
   "source": [
    "## License\n",
    "\n",
    "This software component is licensed by ST under BSD-3-Clause license,\n",
    "the \"License\"; \n",
    "\n",
    "You may not use this file except in compliance with the\n",
    "License. \n",
    "\n",
    "You may obtain a copy of the License at: https://opensource.org/licenses/BSD-3-Clause\n",
    "\n",
    "Copyright (c) 2023 STMicroelectronics. All rights reserved"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cb6e84b0",
   "metadata": {},
   "source": [
    "# Creating a tltb model for Mobilenet v2 using Pytorch\n",
    "This notebook presents scripts to take a pretarained Mobilenet v2 model and export it as onnx and a tao model. \n",
    "In the following sections you can see:\n",
    "* [how can you start from your own model in PyTorch, load the weights to it from a public repo and export it as an ONNX model](#head-1), \n",
    "* [convert it to tltb model which you can to bring to TAO framework](#head-2). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e686a622",
   "metadata": {},
   "source": [
    "#### (Optional) A. Set proxy variables if working behind corporate proxies.\n",
    "\n",
    "The following section sets the proxies and ssl verification flag when the users are working behind the proxies. This setup is necessary to be able to communicate with internet.\n",
    "\n",
    "Replace the `userName`, `password`, and `proxy_port` with your correct username, password and proxy port."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73bd8073",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set proxies\n",
    "import os\n",
    "# os.environ[\"http_proxy\"]='http://userName:password.@example.com:proxy_port'\n",
    "# os.environ[\"https_proxy\"] = 'http://userName:password.@example.com:proxy_port'\n",
    "# os.environ[\"NO_SSL_VERIFY\"]=\"1\"\n",
    "# os.environ[\"SSL_VERIFY\"]=\"False\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1a7ee391",
   "metadata": {},
   "source": [
    "## Creating a tltb model for mobilenetv2 using Pytorch <a class=\"anchor\" id=\"head-1\"></a>\n",
    "Installing all the dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2910921",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install efficientnet-pytorch==0.6.3 \n",
    "!pip install segmentation-models-pytorch==0.2.1 \n",
    "!pip install timm torch==1.10.1\n",
    "!pip install torchvision==0.11.2\n",
    "!pip install tqdm==4.64.0\n",
    "!pip install netron\n",
    "!pip install tensorflow==2.9.1\n",
    "!pip install nvidia-tao-byom\n",
    "!pip install numpy==1.23.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ee642b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the required dependecies\n",
    "import torch\n",
    "import sys\n",
    "sys.path.append('./mobilenetv2_utils/')\n",
    "from mobilenetv2 import mobilenetv2\n",
    "\n",
    "# variables to be used during onnx model creation\n",
    "model_name = \"mobilenet_v2\"\n",
    "ALPHA = 0.5\n",
    "resolution = 128\n",
    "pretrained_weights = \"./mobilenetv2_utils/mobilenetv2_0.5-eaa6f9ad.pth\"\n",
    "\n",
    "''' \n",
    " Create a model using the mobilenetv2.py file and load the weights. \n",
    " The script \"mobilenetv2.py\" and weights \"mobilenetv2_0.5-eaa6f9ad.pth\"\n",
    " can be obtained from https://github.com/d-li14/mobilenetv2.pytorch.\n",
    "'''\n",
    "model = mobilenetv2(width_mult=ALPHA)\n",
    "model.load_state_dict(torch.load(pretrained_weights))\n",
    "\n",
    "'''\n",
    "alternatively one can create the model using torch vision api directly as below\n",
    "model = torch.hub.load('pytorch/vision:v0.10.0', torch_model_name, width_mult=ALPHA, pretrained=pretrained_weights)\n",
    "'''\n",
    "\n",
    "\n",
    "# export the model as onnx\n",
    "dummy_input = torch.randn(1, 3, resolution, resolution, requires_grad=True) \n",
    "\n",
    "# this is needed as inpot to onnx.export\n",
    "export_name = f\"{model_name}_{resolution}_{str(ALPHA).replace('.', '_')}.onnx\"\n",
    "\n",
    "# exporting a torch model as .onnx\n",
    "torch.onnx.export(model, dummy_input, export_name, input_names=[\"input_1\"], verbose=False,\n",
    "                  training=\"TrainingMode.TRAIN\", do_constant_folding=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "76f21d4f",
   "metadata": {},
   "source": [
    "## 1.1: View the network using netron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "100d200d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import netron\n",
    "netron.start(export_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "640ed910",
   "metadata": {},
   "source": [
    "# 2. Convert .onnx to .tltb with TAO-Byom<a class=\"anchor\" id=\"head-2\"></a>\n",
    "This step is to create the .tltb file from the .onnx file. .tltb is an internal format of TAO and this step is needed to start training using TAO. This step will create a byom_model directory with the file and complementary information for this in it.\n",
    "\n",
    "If we wish to fine tune the pretrained model (as is the case) on a different dataset through TAO Toolkit, we must remove the classification head for ImageNet. Hence, the final converted TAO model should only contain layers up to the penultimate layer, which is a layer before the average pooling. In this case, the node name is `463`. Adding `-p 463` argument, removes the head of the model while creatig the `.tltb` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7667fb8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env TF_ENABLE_ONEDNN_OPTS=1\n",
    "# creating the .tltb file\n",
    "!tao_byom -m {export_name} -r ../pretrained_mobilenetv2/mobilenetv2_128_0_5 -n mobilenetv2_128_0_5 -k nvidia_tlt -p 463"
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}