notebooks/4_model_onnx_conversion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import onnx\n",
    "\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_full_model = \"../models/full_model.pt\"\n",
    "path_converted_model = \"../models/full_model.onnx\"\n",
    "\n",
    "path_dataset = \"../datasets/dataset.h5py\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "h5f = h5py.File(path_dataset, 'r')\n",
    "\n",
    "X_train = np.array(h5f['X_train'])\n",
    "X_test = np.array(h5f['X_test'])\n",
    "y_train = np.array(h5f['y_train'])\n",
    "y_test = np.array(h5f['y_test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The architecture was copied from the previous notebook [model_creation.ipynb](./model_creation.ipynb). **torch.load()** need the CNN object, so we redifine it here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(CNN, self).__init__()\n",
    "\n",
    "        # Features detector\n",
    "        self.features1 = nn.Sequential(\n",
    "            \n",
    "            # Hidden layer 1\n",
    "            nn.Conv2d(3, 16, kernel_size=3), nn.ReLU(),\n",
    "            nn.BatchNorm2d(16),\n",
    "\n",
    "            nn.MaxPool2d((2, 2)),\n",
    "            nn.Dropout(0.35)\n",
    "        )\n",
    "\n",
    "        # Features detector\n",
    "        self.features2 = nn.Sequential(\n",
    "            \n",
    "            # Hidden layer 3\n",
    "            nn.Conv2d(16, 32, kernel_size=3), nn.ReLU(),\n",
    "            nn.BatchNorm2d(32),\n",
    "\n",
    "            nn.MaxPool2d((2, 2)),\n",
    "            nn.Dropout(0.35)\n",
    "        )\n",
    "\n",
    "        # Features detector\n",
    "        self.features3 = nn.Sequential(\n",
    "            \n",
    "            # Hidden layer 3\n",
    "            nn.Conv2d(32, 64, kernel_size=3), nn.ReLU(),\n",
    "            nn.BatchNorm2d(64),\n",
    "\n",
    "            nn.MaxPool2d((2, 2)),\n",
    "            nn.Dropout(0.35)\n",
    "        )\n",
    "\n",
    "        # Classifier\n",
    "        self.classifier = nn.Sequential(\n",
    "            \n",
    "            nn.Linear(64*13*13, 512), nn.ReLU(),\n",
    "            nn.BatchNorm1d(512),\n",
    "            nn.Dropout(0.50),\n",
    "\n",
    "            # Output layer\n",
    "            nn.Linear(512, 108)\n",
    "        )\n",
    "\n",
    "    def forward(self, X):\n",
    "\n",
    "        # Features\n",
    "        X = self.features1(X)\n",
    "        X = self.features2(X)\n",
    "        X = self.features3(X)\n",
    "\n",
    "        # print(X.shape)\n",
    "\n",
    "        # Classifier\n",
    "        X = X.view(-1, 64*13*13)\n",
    "        #X = X.view(X.size(0), -1)\n",
    "        X = self.classifier(X)\n",
    "\n",
    "        return F.log_softmax(X, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CNN(\n",
       "  (features1): Sequential(\n",
       "    (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (3): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
       "    (4): Dropout(p=0.35, inplace=False)\n",
       "  )\n",
       "  (features2): Sequential(\n",
       "    (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (3): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
       "    (4): Dropout(p=0.35, inplace=False)\n",
       "  )\n",
       "  (features3): Sequential(\n",
       "    (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (3): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)\n",
       "    (4): Dropout(p=0.35, inplace=False)\n",
       "  )\n",
       "  (classifier): Sequential(\n",
       "    (0): Linear(in_features=10816, out_features=512, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (3): Dropout(p=0.5, inplace=False)\n",
       "    (4): Linear(in_features=512, out_features=128, bias=True)\n",
       "    (5): ReLU()\n",
       "    (6): Linear(in_features=128, out_features=6, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = torch.load(path_full_model)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert to ONNX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set batch size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is also possible to not set a fixed batch_size, to be able to used multiple batch sizes once the model is in tensorrt: https://github.com/onnx/keras-onnx/issues/605"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = torch.rand(BATCH_SIZE, 3, 120, 120).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.onnx.export(\n",
    "    model,\n",
    "    sample,\n",
    "    path_converted_model,\n",
    "    input_names=[\"input\"],\n",
    "    output_names=[\"output\"],\n",
    "    export_params=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that the model converted fine\n",
    "onnx_model = onnx.load(path_converted_model)\n",
    "onnx.checker.check_model(onnx_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}