pytorch2onnx.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "strong-prior",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch import nn\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vital-thompson",
   "metadata": {},
   "source": [
    "# EfficientNet-B0 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "electric-sector",
   "metadata": {
    "code_folding": [
     1
    ]
   },
   "outputs": [],
   "source": [
    "from efficientnet_pytorch import EfficientNet\n",
    "class MyModel(nn.Module):\n",
    "    def __init__(self, model_name: str):\n",
    "        super(MyModel,self).__init__()\n",
    "\n",
    "        # self.model = torch.hub.load('pytorch/vision:v0.9.0', 'inception_v3', pretrained=True)\n",
    "        self.model = EfficientNet.from_pretrained(model_name)\n",
    "        self.model.set_swish(memory_efficient=False)\n",
    "        self.linear1 = torch.nn.Linear(1000, 1)\n",
    "        self.activation3 = torch.nn.Sigmoid()\n",
    "\n",
    "    def forward(self,x):\n",
    "        answ = self.model(x)\n",
    "        #print(answ.shape)\n",
    "        #answ = self.encoder(answ)\n",
    "        #print(answ.shape)\n",
    "        answ = self.activation3(self.linear1(answ))\n",
    "\n",
    "        return answ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "romance-neighbor",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "transform_efnet=transforms.Compose([transforms.Resize((150,150)),\n",
    "                              transforms.ToTensor()\n",
    "                              ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "sorted-yield",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def init_ef_model(model_name, model_path):\n",
    "    device = torch.device('cpu')\n",
    "    model = MyModel(model_name)\n",
    "    model.load_state_dict(torch.load(model_path, map_location=device))\n",
    "    model.eval();\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "accurate-hawaiian",
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def get_res_efnet(image_path: str, model: MyModel):\n",
    "    image = Image.open(image_path)\n",
    "    image = transform_efnet(image).unsqueeze(0)\n",
    "\n",
    "    # model.cuda()\n",
    "    # Set layers such as dropout and batchnorm in evaluation mode\n",
    "    # model.eval();\n",
    "\n",
    "    # Get the 1000-dimensional model output\n",
    "    with torch.no_grad():\n",
    "        out = model(image)\n",
    "    # print(out.detach().numpy()[0][0])    \n",
    "    return out.detach().numpy()[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "seeing-sperm",
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_input = torch.randn(1, 3, 150, 150, device='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "polish-tucson",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained weights for efficientnet-b0\n"
     ]
    }
   ],
   "source": [
    "model_b0_name = 'efficientnet-b0'\n",
    "model_b0_path = 'PretrainedModels/GenderEstimation/Ours/ef_b0.pth'\n",
    "b0_model = init_ef_model(model_b0_name, model_b0_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "banned-difficulty",
   "metadata": {},
   "source": [
    "# Export model to ONNX format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "weighted-assignment",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.onnx.export(b0_model, dummy_input, \"efnet_b0.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "inclusive-finish",
   "metadata": {},
   "source": [
    "# Check model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "round-joining",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import onnx\n",
    "\n",
    "# Load the ONNX model\n",
    "model = onnx.load(\"efnet_b0.onnx\")\n",
    "\n",
    "# Check that the IR is well formed\n",
    "onnx.checker.check_model(model)\n",
    "\n",
    "# Print a human readable representation of the graph\n",
    "# print(onnx.helper.printable_graph(model.graph))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ultimate-vegetation",
   "metadata": {},
   "source": [
    "# Usage onnx model and checking results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "bridal-throw",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exported model has been tested with ONNXRuntime, and the result looks good!\n",
      "Torch prediction: 0.94459\n",
      "ONNX prediction: 0.94459\n"
     ]
    }
   ],
   "source": [
    "import onnxruntime\n",
    "\n",
    "ort_session = onnxruntime.InferenceSession(\"efnet_b0.onnx\")\n",
    "\n",
    "def to_numpy(tensor):\n",
    "    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n",
    "\n",
    "# compute ONNX Runtime output prediction\n",
    "ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy_input)}\n",
    "ort_outs = ort_session.run(None, ort_inputs)\n",
    "\n",
    "# compute pytorch prediction\n",
    "torch_out = b0_model(dummy_input)\n",
    "\n",
    "# compare ONNX Runtime and PyTorch results\n",
    "np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)\n",
    "print(\"Exported model has been tested with ONNXRuntime, and the result looks good!\")\n",
    "\n",
    "torch_res = torch_out.detach().numpy()[0][0]\n",
    "onnx_res = ort_outs[0][0][0]\n",
    "print(\"Torch prediction: %0.5f\" % torch_res + \"\\nONNX prediction: %0.5f\" % onnx_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "electoral-history",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "specific-conflict",
   "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.7.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}