ch-05/lightning_to_onnx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "69075943",
   "metadata": {},
   "source": [
    "### Save the Lightning MNIST model to ONNX format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7c351760",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from torchvision import transforms\n",
    "\n",
    "# from here we get MNIST dataset\n",
    "from torchvision.datasets import MNIST\n",
    "\n",
    "from light_mnist_cnn import LitMNISTCNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "05a3cd9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# where we're storing the downloaded dataset\n",
    "PATH_DATASETS = \".\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5783970",
   "metadata": {},
   "source": [
    "#### Reload the model from a checkpoint and save in ONNX format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5924c78b",
   "metadata": {},
   "outputs": [],
   "source": [
    "CKPT_PATH = \"/home/datascience/pytorch-on-oci/ch-04/checkpoint_mnist/best.ckpt\"\n",
    "\n",
    "model = LitMNISTCNN.load_from_checkpoint(CKPT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "873dc5b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LitMNISTCNN(\n",
       "  (model): Sequential(\n",
       "    (0): Conv2d(1, 32, kernel_size=(5, 5), stride=(1, 1))\n",
       "    (1): ReLU()\n",
       "    (2): Conv2d(32, 32, kernel_size=(5, 5), stride=(1, 1))\n",
       "    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (4): ReLU()\n",
       "    (5): Dropout(p=0.5, inplace=False)\n",
       "    (6): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))\n",
       "    (7): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (8): ReLU()\n",
       "    (9): Dropout(p=0.5, inplace=False)\n",
       "    (10): Flatten(start_dim=1, end_dim=-1)\n",
       "    (11): Linear(in_features=576, out_features=256, bias=True)\n",
       "    (12): ReLU()\n",
       "    (13): Dropout(p=0.1, inplace=False)\n",
       "    (14): Linear(in_features=256, out_features=10, bias=True)\n",
       "  )\n",
       "  (val_accuracy): Accuracy()\n",
       "  (test_accuracy): Accuracy()\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get the summary of CNN architecture\n",
    "\n",
    "# I have moved the Lightmodule defining the CNN to n external py file (see import)\n",
    "\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a32612ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we need an input sample to save in onnx format\n",
    "\n",
    "# the definition of transforms over images\n",
    "img_transforms = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        # normalization is clarified here\n",
    "        # https://discuss.pytorch.org/t/normalization-in-the-mnist-example/457\n",
    "        transforms.Normalize((0.1307,), (0.3081,)),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# we take an input image from the dataset\n",
    "# when we load the dataset we apply transforms as expected from the model\n",
    "mnist_test = MNIST(\".\", train=False, download=True, transform=img_transforms)\n",
    "\n",
    "# index of the image for the test\n",
    "INDEX = 5\n",
    "\n",
    "# take a sample\n",
    "img_tensor, label = mnist_test[INDEX]\n",
    "\n",
    "# make it a batch\n",
    "input_batch = img_tensor.unsqueeze(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ba00278",
   "metadata": {},
   "source": [
    "#### Save in ONNX format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e82c3ecd",
   "metadata": {},
   "outputs": [],
   "source": [
    "ONNX_FILE = \"light_cnn.onnx\"\n",
    "\n",
    "model.to_onnx(ONNX_FILE, input_batch, export_params=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59dda411",
   "metadata": {},
   "source": [
    "#### Test using onnxruntime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "60476521",
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "import onnxruntime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "5e0d506a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model is valid!\n"
     ]
    }
   ],
   "source": [
    "# check the model (takes some seconds)\n",
    "onnx_model = onnx.load(ONNX_FILE)\n",
    "\n",
    "try:\n",
    "    onnx.checker.check_model(onnx_model, full_check=True)\n",
    "except onnx.checker.ValidationError as e:\n",
    "    print(\"The model is invalid: %s\" % e)\n",
    "else:\n",
    "    print(\"The model is valid!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3c3976e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# only at setup\n",
    "ort_session = onnxruntime.InferenceSession(ONNX_FILE)\n",
    "\n",
    "input_name = ort_session.get_inputs()[0].name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "436688e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicted label is: 1, expected label is: 1\n"
     ]
    }
   ],
   "source": [
    "# as input you need to pass a np array, not torch tensor\n",
    "# otherwise it wouldn't be... independent from the framework !\n",
    "ort_inputs = {input_name: input_batch.numpy()}\n",
    "\n",
    "out_class = np.argmax(ort_session.run(None, ort_inputs))\n",
    "\n",
    "print()\n",
    "print(f\"Predicted label is: {out_class}, expected label is: {label}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b71fda5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:computervision_p37_gpu_v1]",
   "language": "python",
   "name": "conda-env-computervision_p37_gpu_v1-py"
  },
  "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.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}