tutorials/conversion_onnx/pytorch2onnx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img width=\"150\" alt=\"Logo_ER10\" src=\"https://user-images.githubusercontent.com/3244249/151994514-b584b984-a148-4ade-80ee-0f88b0aefa45.png\">\n",
    "\n",
    "## Pytorch and Pytorch-lightning to ONNX conversion\n",
    "\n",
    "This notebook shows how to convert your trained Pytorch model or Pytorch-lightning model to ONNX, the generic format supported by DIANNA. <br>\n",
    "It is based on tutorial at https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.init as init\n",
    "from torch.utils import model_zoo\n",
    "import onnx\n",
    "import onnxruntime as ort"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an example model using pytorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SuperResolutionNet(nn.Module):\n",
    "    def __init__(self, upscale_factor, inplace=False):\n",
    "        super(SuperResolutionNet, self).__init__()\n",
    "\n",
    "        self.relu = nn.ReLU(inplace=inplace)\n",
    "        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))\n",
    "        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))\n",
    "        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))\n",
    "        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))\n",
    "        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)\n",
    "\n",
    "        self._initialize_weights()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.conv1(x))\n",
    "        x = self.relu(self.conv2(x))\n",
    "        x = self.relu(self.conv3(x))\n",
    "        x = self.pixel_shuffle(self.conv4(x))\n",
    "        return x\n",
    "\n",
    "    def _initialize_weights(self):\n",
    "        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))\n",
    "        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))\n",
    "        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))\n",
    "        init.orthogonal_(self.conv4.weight)\n",
    "\n",
    "# Create the super-resolution model by using the above model definition.\n",
    "torch_model = SuperResolutionNet(upscale_factor=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load pre-determined weights and applied them to the example model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SuperResolutionNet(\n",
       "  (relu): ReLU()\n",
       "  (conv1): Conv2d(1, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv3): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv4): Conv2d(32, 9, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (pixel_shuffle): PixelShuffle(upscale_factor=3)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get existing weights\n",
    "# Load pretrained model weights\n",
    "model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'\n",
    "\n",
    "# Initialize model with the pretrained weights\n",
    "torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=torch.device('cpu')))\n",
    "\n",
    "# set the model to inference mode\n",
    "torch_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate model on some random input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input to the model\n",
    "x = torch.randn(1, 1, 224, 224, requires_grad=True)\n",
    "# Generate predictions\n",
    "pred = torch_model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export the model in ONNX format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_file = 'pytorch_super_resolution_net.onnx'\n",
    "# Export the model\n",
    "torch.onnx.export(torch_model,               # model being run\n",
    "                  x,                         # model input (or a tuple for multiple inputs)\n",
    "                  onnx_file,                 # where to save the model (can be a file or file-like object)\n",
    "                  export_params=True,        # store the trained parameter weights inside the model file\n",
    "                  opset_version=10,          # the ONNX version to export the model to\n",
    "                  do_constant_folding=True,  # whether to execute constant folding for optimization\n",
    "                  input_names = ['input'],   # the model's input names\n",
    "                  output_names = ['output'], # the model's output names\n",
    "                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes\n",
    "                                'output' : {0 : 'batch_size'}})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify the PyTorch and ONNX predictions match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# verify the ONNX model is valid\n",
    "onnx_model = onnx.load(onnx_file)\n",
    "onnx.checker.check_model(onnx_model)\n",
    "\n",
    "# get ONNX predictions\n",
    "sess = ort.InferenceSession(onnx_file)\n",
    "input_name = sess.get_inputs()[0].name\n",
    "output_name = sess.get_outputs()[0].name\n",
    "\n",
    "onnx_input = {input_name: x.detach().numpy().astype(np.float32)}\n",
    "pred_onnx = sess.run([output_name], onnx_input)[0]\n",
    "\n",
    "# compare to Pytorch predictions\n",
    "np.allclose(pred.detach().numpy(), pred_onnx, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your model is generated using Pytorch-lightning, the conversion is similar as above. <br>\n",
    "\n",
    "Below it shows how to convert your trained Pytorch-lightning model to ONNX.<br>\n",
    "It is based on tutorial at https://pytorch-lightning.readthedocs.io/en/latest/deploy/production_advanced.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pytorch-lightning\n",
    "import pytorch_lightning as pl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an example model using pytorch-lightning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SuperResolutionNetLightning(pl.LightningModule):\n",
    "    def __init__(self, upscale_factor, inplace=False):\n",
    "        super(SuperResolutionNetLightning, self).__init__()\n",
    "\n",
    "        self.relu = nn.ReLU(inplace=inplace)\n",
    "        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))\n",
    "        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))\n",
    "        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))\n",
    "        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))\n",
    "        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)\n",
    "\n",
    "        self._initialize_weights()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.conv1(x))\n",
    "        x = self.relu(self.conv2(x))\n",
    "        x = self.relu(self.conv3(x))\n",
    "        x = self.pixel_shuffle(self.conv4(x))\n",
    "        return x\n",
    "\n",
    "    def _initialize_weights(self):\n",
    "        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))\n",
    "        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))\n",
    "        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))\n",
    "        init.orthogonal_(self.conv4.weight)\n",
    "\n",
    "# Create the super-resolution model by using the above model definition.\n",
    "torch_lightning_model = SuperResolutionNetLightning(upscale_factor=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load pre-determined weights and applied them to the example model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SuperResolutionNetLightning(\n",
       "  (relu): ReLU()\n",
       "  (conv1): Conv2d(1, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv3): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv4): Conv2d(32, 9, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (pixel_shuffle): PixelShuffle(upscale_factor=3)\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize model with the pretrained weights\n",
    "torch_lightning_model.load_state_dict(model_zoo.load_url(model_url, map_location=torch.device('cpu')))\n",
    "\n",
    "# set the model to inference mode\n",
    "torch_lightning_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate model on some random input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input to the model\n",
    "x = torch.randn(1, 1, 224, 224, requires_grad=True)\n",
    "# Generate predictions\n",
    "pred = torch_lightning_model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Export the model in ONNX format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export to onnx format\n",
    "onnx_file = 'pytorch_lightning_super_resolution_net.onnx'\n",
    "torch_lightning_model.to_onnx(file_path = onnx_file, input_sample = x, export_params=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify the PyTorch and ONNX predictions match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# verify the ONNX model is valid\n",
    "onnx_model = onnx.load(onnx_file)\n",
    "onnx.checker.check_model(onnx_model)\n",
    "\n",
    "# get ONNX predictions\n",
    "sess = ort.InferenceSession(onnx_file)\n",
    "input_name = sess.get_inputs()[0].name\n",
    "output_name = sess.get_outputs()[0].name\n",
    "\n",
    "onnx_input = {input_name: x.detach().numpy().astype(np.float32)}\n",
    "pred_onnx = sess.run([output_name], onnx_input)[0]\n",
    "\n",
    "# compare to Pytorch predictions\n",
    "np.allclose(pred.detach().numpy(), pred_onnx, atol=1e-5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.8 ('torch')",
   "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.10.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f74811edbe99894b2f930b63702daebe3ce5897f538d47d6f6827e4475af2be0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}