onnx.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:56.630555Z",
     "start_time": "2020-12-23T04:49:56.217495Z"
    }
   },
   "outputs": [],
   "source": [
    "# 필요한 import문\n",
    "import io\n",
    "import numpy as np\n",
    "\n",
    "from torch import nn\n",
    "import torch.utils.model_zoo as model_zoo\n",
    "import torch.onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:56.637544Z",
     "start_time": "2020-12-23T04:49:56.631545Z"
    }
   },
   "outputs": [],
   "source": [
    "# PyTorch에서 구현된 초해상도 모델\n",
    "import torch.nn as nn\n",
    "import torch.nn.init as init\n",
    "\n",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:56.657579Z",
     "start_time": "2020-12-23T04:49:56.639545Z"
    }
   },
   "outputs": [],
   "source": [
    "torch_model = SuperResolutionNet(upscale_factor=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:56.661582Z",
     "start_time": "2020-12-23T04:49:56.659579Z"
    }
   },
   "outputs": [],
   "source": [
    "# 미리 학습된 가중치를 읽어옵니다\n",
    "model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'\n",
    "batch_size = 1    # 임의의 수"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:57.898267Z",
     "start_time": "2020-12-23T04:49:56.663580Z"
    }
   },
   "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": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 모델을 미리 학습된 가중치로 초기화합니다\n",
    "map_location = lambda storage, loc: storage\n",
    "if torch.cuda.is_available():\n",
    "    map_location = None\n",
    "torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))\n",
    "\n",
    "# 모델을 추론 모드로 전환합니다\n",
    "torch_model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:57.904265Z",
     "start_time": "2020-12-23T04:49:57.899266Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          ...,\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "          [1., 1., 1.,  ..., 1., 1., 1.]]]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 모델에 대한 입력값\n",
    "x = torch.ones(1,1,224,224)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:57.955285Z",
     "start_time": "2020-12-23T04:49:57.905265Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0.7651, 0.8741, 0.9048,  ..., 0.8559, 0.7837, 0.7129],\n",
       "          [0.8241, 1.0140, 1.1038,  ..., 1.1016, 1.0106, 0.7954],\n",
       "          [0.8494, 1.0746, 1.1395,  ..., 1.1076, 1.0816, 0.8434],\n",
       "          ...,\n",
       "          [0.9251, 1.0277, 1.0539,  ..., 1.0218, 1.0811, 0.9037],\n",
       "          [0.8515, 0.9133, 0.9483,  ..., 0.9819, 1.0057, 0.8684],\n",
       "          [0.7569, 0.7720, 0.7775,  ..., 0.7539, 0.8017, 0.7727]]]],\n",
       "       grad_fn=<UnsafeViewBackward>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 모델의 출력\n",
    "torch_out = torch_model(x)\n",
    "torch_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.071265Z",
     "start_time": "2020-12-23T04:49:57.957267Z"
    }
   },
   "outputs": [],
   "source": [
    "# 모델 변환\n",
    "torch.onnx.export(torch_model,               # 실행될 모델\n",
    "                  x,                         # 모델 입력값 (튜플 또는 여러 입력값들도 가능)\n",
    "                  \"super_resolution.onnx\",   # 모델 저장 경로 (파일 또는 파일과 유사한 객체 모두 가능)\n",
    "                  export_params=True,        # 모델 파일 안에 학습된 모델 가중치를 저장할지의 여부\n",
    "                  opset_version=10,          # 모델을 변환할 때 사용할 ONNX 버전\n",
    "                  do_constant_folding=True,  # 최적하시 상수폴딩을 사용할지의 여부\n",
    "                  input_names = ['input'],   # 모델의 입력값을 가리키는 이름\n",
    "                  output_names = ['output'], # 모델의 출력값을 가리키는 이름\n",
    "                  dynamic_axes={'input' : {0 : 'batch_size'},    # 가변적인 길이를 가진 차원\n",
    "                                'output' : {0 : 'batch_size'}})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.120264Z",
     "start_time": "2020-12-23T04:49:58.073265Z"
    }
   },
   "outputs": [],
   "source": [
    "import onnx\n",
    "\n",
    "onnx_model = onnx.load(\"super_resolution.onnx\")\n",
    "onnx.checker.check_model(onnx_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.153263Z",
     "start_time": "2020-12-23T04:49:58.121265Z"
    }
   },
   "outputs": [],
   "source": [
    "import onnxruntime\n",
    "\n",
    "ort_session = onnxruntime.InferenceSession(\"super_resolution.onnx\")\n",
    "\n",
    "def to_numpy(tensor):\n",
    "    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.176267Z",
     "start_time": "2020-12-23T04:49:58.154264Z"
    }
   },
   "outputs": [],
   "source": [
    "# ONNX 런타임에서 계산된 결과값\n",
    "ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}\n",
    "\n",
    "ort_outs = ort_session.run(None, ort_inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.183264Z",
     "start_time": "2020-12-23T04:49:58.177266Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[[[0.76511073, 0.8740847 , 0.90481126, ..., 0.85588545,\n",
       "           0.7836736 , 0.7129    ],\n",
       "          [0.82408434, 1.0140268 , 1.1037575 , ..., 1.1016412 ,\n",
       "           1.0106119 , 0.79541755],\n",
       "          [0.84935516, 1.0745931 , 1.1394703 , ..., 1.1076125 ,\n",
       "           1.0815643 , 0.843422  ],\n",
       "          ...,\n",
       "          [0.92509943, 1.0276617 , 1.0539078 , ..., 1.0218285 ,\n",
       "           1.081129  , 0.9036602 ],\n",
       "          [0.851483  , 0.91334385, 0.9482872 , ..., 0.9819409 ,\n",
       "           1.005691  , 0.86839116],\n",
       "          [0.7569455 , 0.772007  , 0.77752113, ..., 0.75390404,\n",
       "           0.8017341 , 0.7727227 ]]]], dtype=float32)]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ort_outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.187264Z",
     "start_time": "2020-12-23T04:49:58.184263Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 1, 672, 672)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ort_outs[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.194265Z",
     "start_time": "2020-12-23T04:49:58.189266Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "451584"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "672*672"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.214267Z",
     "start_time": "2020-12-23T04:49:58.196265Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exported model has been tested with ONNXRuntime, and the result looks good!\n"
     ]
    }
   ],
   "source": [
    "# ONNX 런타임과 PyTorch에서 연산된 결과값 비교\n",
    "np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)\n",
    "\n",
    "print(\"Exported model has been tested with ONNXRuntime, and the result looks good!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.222265Z",
     "start_time": "2020-12-23T04:49:58.215264Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[0.7651, 0.8741, 0.9048,  ..., 0.8559, 0.7837, 0.7129],\n",
       "          [0.8241, 1.0140, 1.1038,  ..., 1.1016, 1.0106, 0.7954],\n",
       "          [0.8494, 1.0746, 1.1395,  ..., 1.1076, 1.0816, 0.8434],\n",
       "          ...,\n",
       "          [0.9251, 1.0277, 1.0539,  ..., 1.0218, 1.0811, 0.9037],\n",
       "          [0.8515, 0.9133, 0.9483,  ..., 0.9819, 1.0057, 0.8684],\n",
       "          [0.7569, 0.7720, 0.7775,  ..., 0.7539, 0.8017, 0.7727]]]],\n",
       "       grad_fn=<UnsafeViewBackward>)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-23T04:49:58.227269Z",
     "start_time": "2020-12-23T04:49:58.223265Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[[[0.76511073, 0.8740847 , 0.90481126, ..., 0.85588545,\n",
       "           0.7836736 , 0.7129    ],\n",
       "          [0.82408434, 1.0140268 , 1.1037575 , ..., 1.1016412 ,\n",
       "           1.0106119 , 0.79541755],\n",
       "          [0.84935516, 1.0745931 , 1.1394703 , ..., 1.1076125 ,\n",
       "           1.0815643 , 0.843422  ],\n",
       "          ...,\n",
       "          [0.92509943, 1.0276617 , 1.0539078 , ..., 1.0218285 ,\n",
       "           1.081129  , 0.9036602 ],\n",
       "          [0.851483  , 0.91334385, 0.9482872 , ..., 0.9819409 ,\n",
       "           1.005691  , 0.86839116],\n",
       "          [0.7569455 , 0.772007  , 0.77752113, ..., 0.75390404,\n",
       "           0.8017341 , 0.7727227 ]]]], dtype=float32)]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ort_outs"
   ]
  }
 ],
 "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.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "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
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}