classification_tf2/byom_converters/byom_converter_resnet18.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d185031a",
   "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 resnet18 using Pytorch\n",
    "This notebook presents scripts to create a pretarained resnet18 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 with pretrained weights and export it as an ONNX model](#head-1), \n",
    "* [convert the exported onnx model into a .tltb model which you can to bring to TAO framework](#head-2). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eed7f01",
   "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": "a87a7b35",
   "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\"\n",
    "\n",
    "# to create an unverified context\n",
    "import ssl\n",
    "ssl._create_default_https_context = ssl._create_unverified_context "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ce751b93",
   "metadata": {},
   "source": [
    "## 1. Create the onnx model. <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3d203f5c",
   "metadata": {},
   "source": [
    "### Import dependencies"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5a01ffa9",
   "metadata": {},
   "source": [
    "Installing all the dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ebfd3e4",
   "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2f774e44",
   "metadata": {},
   "source": [
    "### Download Resnet18 weights\n",
    "Following we download the weights for resnet18 model from `torch.hub`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f9219cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch_model_name = 'resnet18'\n",
    "resolution = 128\n",
    "pretrained_weights = True\n",
    "\n",
    "model = torch.hub.load('pytorch/vision:v0.10.0', torch_model_name,  pretrained=pretrained_weights)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "de0c50de",
   "metadata": {},
   "source": [
    "### Exporting the generated model as .onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c953a4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_ = torch.randn(1, 3, resolution, resolution, requires_grad=True)\n",
    "export_name = torch_model_name + '_' + str(resolution) + '.onnx'\n",
    "torch.onnx.export(model, 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 `188`. Adding `-p 188` 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_resnet18/resnet18_128 -n resnet18_128 -k nvidia_tlt -p 188"
   ]
  }
 ],
 "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
}