{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Converting PyTorch to ONNX" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/gautham/.local/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n" ] } ], "source": [ "import torch\n", "from torch import nn\n", "from torch.nn import functional as F" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Defining the model" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "class DoubleConv(nn.Module):\n", " def __init__(self, in_channels, out_channels):\n", " super().__init__()\n", " self.conv = nn.Sequential(\n", " nn.Conv2d(in_channels, out_channels, 3, padding=1),\n", " nn.BatchNorm2d(out_channels),\n", " nn.ReLU(inplace=True),\n", " nn.Conv2d(out_channels, out_channels, 3, padding=1),\n", " nn.BatchNorm2d(out_channels),\n", " nn.ReLU(inplace=True)\n", " )\n", " def forward(self, x):\n", " return self.conv(x)\n", "\n", "class Down(nn.Module):\n", " def __init__(self, in_channels, out_channels):\n", " super().__init__()\n", " self.down = nn.Sequential(\n", " nn.MaxPool2d(2),\n", " DoubleConv(in_channels, out_channels)\n", " )\n", " def forward(self, x):\n", " return self.down(x)\n", "\n", "class Up(nn.Module):\n", " def __init__(self, in_channels, out_channels, bilinear=False):\n", " super().__init__()\n", " if bilinear:\n", " self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n", " nn.Conv2d(in_channels, in_channels // 2, 1))\n", " else:\n", " self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, 2, stride=2)\n", " \n", " self.conv = DoubleConv(in_channels, out_channels)\n", "\n", " def forward(self, x1, x2):\n", " x1 = self.up(x1)\n", " diffY = x2.size()[2] - x1.size()[2]\n", " diffX = x2.size()[3] - x1.size()[3]\n", " x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])\n", " x = torch.cat([x2, x1], dim=1)\n", " return self.conv(x)\n", "\n", "class OutConv(nn.Module):\n", " def __init__(self, in_channels, out_channels):\n", " super().__init__()\n", " self.conv = nn.Conv2d(in_channels, out_channels, 1)\n", " self.sigmoid = nn.Sigmoid()\n", "\n", " def forward(self, x):\n", " return self.sigmoid(self.conv(x))\n", "\n", "class UNet(nn.Module):\n", " def __init__(self, n_channels, n_classes):\n", " super().__init__()\n", " self.inc = DoubleConv(n_channels, 64)\n", " self.down1 = Down(64, 128)\n", " self.down2 = Down(128, 256)\n", " self.down3 = Down(256, 512)\n", " self.down4 = Down(512, 1024)\n", " self.up1 = Up(1024, 512)\n", " self.up2 = Up(512, 256)\n", " self.up3 = Up(256, 128)\n", " self.up4 = Up(128, 64)\n", " self.outc = OutConv(64, n_classes)\n", "\n", " def forward(self, x):\n", " x1 = self.inc(x)\n", " x2 = self.down1(x1)\n", " x3 = self.down2(x2)\n", " x4 = self.down3(x3)\n", " x5 = self.down4(x4)\n", " x = self.up1(x5, x4)\n", " x = self.up2(x4, x3)\n", " x = self.up3(x3, x2)\n", " x = self.up4(x2, x1)\n", " logits = self.outc(x)\n", " return logits" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Loading the model" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "ename": "RuntimeError", "evalue": "Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)", "Cell \u001b[0;32mIn [3], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m model \u001b[39m=\u001b[39m UNet(n_channels\u001b[39m=\u001b[39m\u001b[39m3\u001b[39m, n_classes\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m model\u001b[39m.\u001b[39mload_state_dict(torch\u001b[39m.\u001b[39;49mload(\u001b[39m'\u001b[39;49m\u001b[39m../weights/water_bodies_model.pth\u001b[39;49m\u001b[39m'\u001b[39;49m))\n\u001b[1;32m 3\u001b[0m model\u001b[39m.\u001b[39mto(\u001b[39m'\u001b[39m\u001b[39mcpu\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m 5\u001b[0m dummy_input \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mrandn(\u001b[39m1\u001b[39m, \u001b[39m3\u001b[39m, \u001b[39m256\u001b[39m, \u001b[39m256\u001b[39m)\n", "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/serialization.py:789\u001b[0m, in \u001b[0;36mload\u001b[0;34m(f, map_location, pickle_module, weights_only, **pickle_load_args)\u001b[0m\n\u001b[1;32m 787\u001b[0m \u001b[39mexcept\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m \u001b[39mas\u001b[39;00m e:\n\u001b[1;32m 788\u001b[0m \u001b[39mraise\u001b[39;00m pickle\u001b[39m.\u001b[39mUnpicklingError(UNSAFE_MESSAGE \u001b[39m+\u001b[39m \u001b[39mstr\u001b[39m(e)) \u001b[39mfrom\u001b[39;00m \u001b[39mNone\u001b[39m\n\u001b[0;32m--> 789\u001b[0m \u001b[39mreturn\u001b[39;00m _load(opened_zipfile, map_location, pickle_module, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mpickle_load_args)\n\u001b[1;32m 790\u001b[0m \u001b[39mif\u001b[39;00m weights_only:\n\u001b[1;32m 791\u001b[0m \u001b[39mtry\u001b[39;00m:\n", "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/serialization.py:1131\u001b[0m, in \u001b[0;36m_load\u001b[0;34m(zip_file, map_location, pickle_module, pickle_file, **pickle_load_args)\u001b[0m\n\u001b[1;32m 1129\u001b[0m unpickler \u001b[39m=\u001b[39m UnpicklerWrapper(data_file, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mpickle_load_args)\n\u001b[1;32m 1130\u001b[0m unpickler\u001b[39m.\u001b[39mpersistent_load \u001b[39m=\u001b[39m persistent_load\n\u001b[0;32m-> 1131\u001b[0m result \u001b[39m=\u001b[39m unpickler\u001b[39m.\u001b[39;49mload()\n\u001b[1;32m 1133\u001b[0m torch\u001b[39m.\u001b[39m_utils\u001b[39m.\u001b[39m_validate_loaded_sparse_tensors()\n\u001b[1;32m 1135\u001b[0m \u001b[39mreturn\u001b[39;00m result\n", "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/serialization.py:1101\u001b[0m, in \u001b[0;36m_load..persistent_load\u001b[0;34m(saved_id)\u001b[0m\n\u001b[1;32m 1099\u001b[0m \u001b[39mif\u001b[39;00m key \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m loaded_storages:\n\u001b[1;32m 1100\u001b[0m nbytes \u001b[39m=\u001b[39m numel \u001b[39m*\u001b[39m torch\u001b[39m.\u001b[39m_utils\u001b[39m.\u001b[39m_element_size(dtype)\n\u001b[0;32m-> 1101\u001b[0m load_tensor(dtype, nbytes, key, _maybe_decode_ascii(location))\n\u001b[1;32m 1103\u001b[0m \u001b[39mreturn\u001b[39;00m loaded_storages[key]\n", "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/serialization.py:1083\u001b[0m, in \u001b[0;36m_load..load_tensor\u001b[0;34m(dtype, numel, key, location)\u001b[0m\n\u001b[1;32m 1079\u001b[0m storage \u001b[39m=\u001b[39m zip_file\u001b[39m.\u001b[39mget_storage_from_record(name, numel, torch\u001b[39m.\u001b[39mUntypedStorage)\u001b[39m.\u001b[39mstorage()\u001b[39m.\u001b[39muntyped()\n\u001b[1;32m 1080\u001b[0m \u001b[39m# TODO: Once we decide to break serialization FC, we can\u001b[39;00m\n\u001b[1;32m 1081\u001b[0m \u001b[39m# stop wrapping with TypedStorage\u001b[39;00m\n\u001b[1;32m 1082\u001b[0m loaded_storages[key] \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mstorage\u001b[39m.\u001b[39mTypedStorage(\n\u001b[0;32m-> 1083\u001b[0m wrap_storage\u001b[39m=\u001b[39mrestore_location(storage, location),\n\u001b[1;32m 1084\u001b[0m dtype\u001b[39m=\u001b[39mdtype)\n", "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/serialization.py:215\u001b[0m, in \u001b[0;36mdefault_restore_location\u001b[0;34m(storage, location)\u001b[0m\n\u001b[1;32m 213\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdefault_restore_location\u001b[39m(storage, location):\n\u001b[1;32m 214\u001b[0m \u001b[39mfor\u001b[39;00m _, _, fn \u001b[39min\u001b[39;00m _package_registry:\n\u001b[0;32m--> 215\u001b[0m result \u001b[39m=\u001b[39m fn(storage, location)\n\u001b[1;32m 216\u001b[0m \u001b[39mif\u001b[39;00m result \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m 217\u001b[0m \u001b[39mreturn\u001b[39;00m result\n", "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/serialization.py:182\u001b[0m, in \u001b[0;36m_cuda_deserialize\u001b[0;34m(obj, location)\u001b[0m\n\u001b[1;32m 180\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_cuda_deserialize\u001b[39m(obj, location):\n\u001b[1;32m 181\u001b[0m \u001b[39mif\u001b[39;00m location\u001b[39m.\u001b[39mstartswith(\u001b[39m'\u001b[39m\u001b[39mcuda\u001b[39m\u001b[39m'\u001b[39m):\n\u001b[0;32m--> 182\u001b[0m device \u001b[39m=\u001b[39m validate_cuda_device(location)\n\u001b[1;32m 183\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mgetattr\u001b[39m(obj, \u001b[39m\"\u001b[39m\u001b[39m_torch_load_uninitialized\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39mFalse\u001b[39;00m):\n\u001b[1;32m 184\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mcuda\u001b[39m.\u001b[39mdevice(device):\n", "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/serialization.py:166\u001b[0m, in \u001b[0;36mvalidate_cuda_device\u001b[0;34m(location)\u001b[0m\n\u001b[1;32m 163\u001b[0m device \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mcuda\u001b[39m.\u001b[39m_utils\u001b[39m.\u001b[39m_get_device_index(location, \u001b[39mTrue\u001b[39;00m)\n\u001b[1;32m 165\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m torch\u001b[39m.\u001b[39mcuda\u001b[39m.\u001b[39mis_available():\n\u001b[0;32m--> 166\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\u001b[39m'\u001b[39m\u001b[39mAttempting to deserialize object on a CUDA \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m 167\u001b[0m \u001b[39m'\u001b[39m\u001b[39mdevice but torch.cuda.is_available() is False. \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m 168\u001b[0m \u001b[39m'\u001b[39m\u001b[39mIf you are running on a CPU-only machine, \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m 169\u001b[0m \u001b[39m'\u001b[39m\u001b[39mplease use torch.load with map_location=torch.device(\u001b[39m\u001b[39m\\'\u001b[39;00m\u001b[39mcpu\u001b[39m\u001b[39m\\'\u001b[39;00m\u001b[39m) \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m 170\u001b[0m \u001b[39m'\u001b[39m\u001b[39mto map your storages to the CPU.\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m 171\u001b[0m device_count \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mcuda\u001b[39m.\u001b[39mdevice_count()\n\u001b[1;32m 172\u001b[0m \u001b[39mif\u001b[39;00m device \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m device_count:\n", "\u001b[0;31mRuntimeError\u001b[0m: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU." ] } ], "source": [ "model = UNet(n_channels=3, n_classes=1)\n", "model.load_state_dict(torch.load('../weights/water_bodies_model.pth'))\n", "\n", "dummy_input = torch.randn(1, 3, 256, 256)\n", "\n", "model.eval()\n", "torch_out = model(dummy_input)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Converting to ONNX" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Exported graph: graph(%input : Float(*, 3, 256, 256, strides=[196608, 65536, 256, 1], requires_grad=0, device=cpu),\n", " %up4.up.weight : Float(128, 64, 2, 2, strides=[256, 4, 2, 1], requires_grad=1, device=cpu),\n", " %up4.up.bias : Float(64, strides=[1], requires_grad=1, device=cpu),\n", " %outc.conv.weight : Float(1, 64, 1, 1, strides=[64, 1, 1, 1], requires_grad=1, device=cpu),\n", " %outc.conv.bias : Float(1, strides=[1], requires_grad=1, device=cpu),\n", " %onnx::Conv_226 : Float(64, 3, 3, 3, strides=[27, 9, 3, 1], requires_grad=0, device=cpu),\n", " %onnx::Conv_227 : Float(64, strides=[1], requires_grad=0, device=cpu),\n", " %onnx::Conv_229 : Float(64, 64, 3, 3, strides=[576, 9, 3, 1], requires_grad=0, device=cpu),\n", " %onnx::Conv_230 : Float(64, strides=[1], requires_grad=0, device=cpu),\n", " %onnx::Conv_232 : Float(128, 64, 3, 3, strides=[576, 9, 3, 1], requires_grad=0, device=cpu),\n", " %onnx::Conv_233 : Float(128, strides=[1], requires_grad=0, device=cpu),\n", " %onnx::Conv_235 : Float(128, 128, 3, 3, strides=[1152, 9, 3, 1], requires_grad=0, device=cpu),\n", " %onnx::Conv_236 : Float(128, strides=[1], requires_grad=0, device=cpu),\n", " %onnx::Conv_238 : Float(64, 128, 3, 3, strides=[1152, 9, 3, 1], requires_grad=0, device=cpu),\n", " %onnx::Conv_239 : Float(64, strides=[1], requires_grad=0, device=cpu),\n", " %onnx::Conv_241 : Float(64, 64, 3, 3, strides=[576, 9, 3, 1], requires_grad=0, device=cpu),\n", " %onnx::Conv_242 : Float(64, strides=[1], requires_grad=0, device=cpu)):\n", " %/inc/conv/conv.0/Conv_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], onnx_name=\"/inc/conv/conv.0/Conv\"](%input, %onnx::Conv_226, %onnx::Conv_227), scope: __main__.UNet::/__main__.DoubleConv::inc/torch.nn.modules.container.Sequential::conv/torch.nn.modules.conv.Conv2d::conv.0 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0\n", " %/inc/conv/conv.2/Relu_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Relu[onnx_name=\"/inc/conv/conv.2/Relu\"](%/inc/conv/conv.0/Conv_output_0), scope: __main__.UNet::/__main__.DoubleConv::inc/torch.nn.modules.container.Sequential::conv/torch.nn.modules.activation.ReLU::conv.2 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/functional.py:1453:0\n", " %/inc/conv/conv.3/Conv_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], onnx_name=\"/inc/conv/conv.3/Conv\"](%/inc/conv/conv.2/Relu_output_0, %onnx::Conv_229, %onnx::Conv_230), scope: __main__.UNet::/__main__.DoubleConv::inc/torch.nn.modules.container.Sequential::conv/torch.nn.modules.conv.Conv2d::conv.3 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0\n", " %/inc/conv/conv.5/Relu_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Relu[onnx_name=\"/inc/conv/conv.5/Relu\"](%/inc/conv/conv.3/Conv_output_0), scope: __main__.UNet::/__main__.DoubleConv::inc/torch.nn.modules.container.Sequential::conv/torch.nn.modules.activation.ReLU::conv.5 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/functional.py:1453:0\n", " %/down1/down/down.0/MaxPool_output_0 : Float(*, 64, 128, 128, strides=[1048576, 16384, 128, 1], requires_grad=1, device=cpu) = onnx::MaxPool[ceil_mode=0, kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2], onnx_name=\"/down1/down/down.0/MaxPool\"](%/inc/conv/conv.5/Relu_output_0), scope: __main__.UNet::/__main__.Down::down1/torch.nn.modules.container.Sequential::down/torch.nn.modules.pooling.MaxPool2d::down.0 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/functional.py:780:0\n", " %/down1/down/down.1/conv/conv.0/Conv_output_0 : Float(*, 128, 128, 128, strides=[2097152, 16384, 128, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], onnx_name=\"/down1/down/down.1/conv/conv.0/Conv\"](%/down1/down/down.0/MaxPool_output_0, %onnx::Conv_232, %onnx::Conv_233), scope: __main__.UNet::/__main__.Down::down1/torch.nn.modules.container.Sequential::down/__main__.DoubleConv::down.1/torch.nn.modules.container.Sequential::conv/torch.nn.modules.conv.Conv2d::conv.0 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0\n", " %/down1/down/down.1/conv/conv.2/Relu_output_0 : Float(*, 128, 128, 128, strides=[2097152, 16384, 128, 1], requires_grad=1, device=cpu) = onnx::Relu[onnx_name=\"/down1/down/down.1/conv/conv.2/Relu\"](%/down1/down/down.1/conv/conv.0/Conv_output_0), scope: __main__.UNet::/__main__.Down::down1/torch.nn.modules.container.Sequential::down/__main__.DoubleConv::down.1/torch.nn.modules.container.Sequential::conv/torch.nn.modules.activation.ReLU::conv.2 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/functional.py:1453:0\n", " %/down1/down/down.1/conv/conv.3/Conv_output_0 : Float(*, 128, 128, 128, strides=[2097152, 16384, 128, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], onnx_name=\"/down1/down/down.1/conv/conv.3/Conv\"](%/down1/down/down.1/conv/conv.2/Relu_output_0, %onnx::Conv_235, %onnx::Conv_236), scope: __main__.UNet::/__main__.Down::down1/torch.nn.modules.container.Sequential::down/__main__.DoubleConv::down.1/torch.nn.modules.container.Sequential::conv/torch.nn.modules.conv.Conv2d::conv.3 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0\n", " %/down1/down/down.1/conv/conv.5/Relu_output_0 : Float(*, 128, 128, 128, strides=[2097152, 16384, 128, 1], requires_grad=1, device=cpu) = onnx::Relu[onnx_name=\"/down1/down/down.1/conv/conv.5/Relu\"](%/down1/down/down.1/conv/conv.3/Conv_output_0), scope: __main__.UNet::/__main__.Down::down1/torch.nn.modules.container.Sequential::down/__main__.DoubleConv::down.1/torch.nn.modules.container.Sequential::conv/torch.nn.modules.activation.ReLU::conv.5 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/functional.py:1453:0\n", " %/up4/up/ConvTranspose_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=0, device=cpu) = onnx::ConvTranspose[dilations=[1, 1], group=1, kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2], onnx_name=\"/up4/up/ConvTranspose\"](%/down1/down/down.1/conv/conv.5/Relu_output_0, %up4.up.weight, %up4.up.bias), scope: __main__.UNet::/__main__.Up::up4/torch.nn.modules.conv.ConvTranspose2d::up # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:953:0\n", " %/up4/Shape_output_0 : Long(4, strides=[1], device=cpu) = onnx::Shape[onnx_name=\"/up4/Shape\"](%/inc/conv/conv.5/Relu_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:37:0\n", " %/up4/Constant_output_0 : Long(device=cpu) = onnx::Constant[value={2}, onnx_name=\"/up4/Constant\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:37:0\n", " %/up4/Gather_output_0 : Long(device=cpu) = onnx::Gather[axis=0, onnx_name=\"/up4/Gather\"](%/up4/Shape_output_0, %/up4/Constant_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:37:0\n", " %/up4/Shape_1_output_0 : Long(4, strides=[1], device=cpu) = onnx::Shape[onnx_name=\"/up4/Shape_1\"](%/up4/up/ConvTranspose_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Constant_1_output_0 : Long(device=cpu) = onnx::Constant[value={2}, onnx_name=\"/up4/Constant_1\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Gather_1_output_0 : Long(device=cpu) = onnx::Gather[axis=0, onnx_name=\"/up4/Gather_1\"](%/up4/Shape_1_output_0, %/up4/Constant_1_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Sub_output_0 : Long(requires_grad=0, device=cpu) = onnx::Sub[onnx_name=\"/up4/Sub\"](%/up4/Gather_output_0, %/up4/Gather_1_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Shape_2_output_0 : Long(4, strides=[1], device=cpu) = onnx::Shape[onnx_name=\"/up4/Shape_2\"](%/inc/conv/conv.5/Relu_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Constant_2_output_0 : Long(device=cpu) = onnx::Constant[value={3}, onnx_name=\"/up4/Constant_2\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Gather_2_output_0 : Long(device=cpu) = onnx::Gather[axis=0, onnx_name=\"/up4/Gather_2\"](%/up4/Shape_2_output_0, %/up4/Constant_2_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Shape_3_output_0 : Long(4, strides=[1], device=cpu) = onnx::Shape[onnx_name=\"/up4/Shape_3\"](%/up4/up/ConvTranspose_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Constant_3_output_0 : Long(device=cpu) = onnx::Constant[value={3}, onnx_name=\"/up4/Constant_3\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Gather_3_output_0 : Long(device=cpu) = onnx::Gather[axis=0, onnx_name=\"/up4/Gather_3\"](%/up4/Shape_3_output_0, %/up4/Constant_3_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Sub_1_output_0 : Long(requires_grad=0, device=cpu) = onnx::Sub[onnx_name=\"/up4/Sub_1\"](%/up4/Gather_2_output_0, %/up4/Gather_3_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:38:0\n", " %/up4/Constant_4_output_0 : Long(requires_grad=0, device=cpu) = onnx::Constant[value={2}, onnx_name=\"/up4/Constant_4\"](), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Div_output_0 : Long(device=cpu) = onnx::Div[onnx_name=\"/up4/Div\"](%/up4/Sub_1_output_0, %/up4/Constant_4_output_0), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Cast_output_0 : Long(device=cpu) = onnx::Cast[to=7, onnx_name=\"/up4/Cast\"](%/up4/Div_output_0), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Cast_1_output_0 : Long(requires_grad=0, device=cpu) = onnx::Cast[to=7, onnx_name=\"/up4/Cast_1\"](%/up4/Cast_output_0), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Sub_2_output_0 : Long(requires_grad=0, device=cpu) = onnx::Sub[onnx_name=\"/up4/Sub_2\"](%/up4/Sub_1_output_0, %/up4/Cast_1_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_5_output_0 : Long(requires_grad=0, device=cpu) = onnx::Constant[value={2}, onnx_name=\"/up4/Constant_5\"](), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Div_1_output_0 : Long(device=cpu) = onnx::Div[onnx_name=\"/up4/Div_1\"](%/up4/Sub_output_0, %/up4/Constant_5_output_0), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Cast_2_output_0 : Long(device=cpu) = onnx::Cast[to=7, onnx_name=\"/up4/Cast_2\"](%/up4/Div_1_output_0), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Cast_3_output_0 : Long(requires_grad=0, device=cpu) = onnx::Cast[to=7, onnx_name=\"/up4/Cast_3\"](%/up4/Cast_2_output_0), scope: __main__.UNet::/__main__.Up::up4 # /home/gautham/.local/lib/python3.10/site-packages/torch/_tensor.py:867:0\n", " %/up4/Sub_3_output_0 : Long(requires_grad=0, device=cpu) = onnx::Sub[onnx_name=\"/up4/Sub_3\"](%/up4/Sub_output_0, %/up4/Cast_3_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %onnx::Unsqueeze_175 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze\"](%/up4/Cast_1_output_0, %onnx::Unsqueeze_175), scope: __main__.UNet::/__main__.Up::up4\n", " %onnx::Unsqueeze_177 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_1_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze_1\"](%/up4/Sub_2_output_0, %onnx::Unsqueeze_177), scope: __main__.UNet::/__main__.Up::up4\n", " %onnx::Unsqueeze_179 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_2_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze_2\"](%/up4/Cast_3_output_0, %onnx::Unsqueeze_179), scope: __main__.UNet::/__main__.Up::up4\n", " %onnx::Unsqueeze_181 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_3_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze_3\"](%/up4/Sub_3_output_0, %onnx::Unsqueeze_181), scope: __main__.UNet::/__main__.Up::up4\n", " %/up4/Concat_output_0 : Long(4, strides=[1], device=cpu) = onnx::Concat[axis=0, onnx_name=\"/up4/Concat\"](%/up4/Unsqueeze_output_0, %/up4/Unsqueeze_1_output_0, %/up4/Unsqueeze_2_output_0, %/up4/Unsqueeze_3_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %onnx::Unsqueeze_184 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_4_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze_4\"](%/up4/Cast_1_output_0, %onnx::Unsqueeze_184), scope: __main__.UNet::/__main__.Up::up4\n", " %onnx::Unsqueeze_186 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_5_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze_5\"](%/up4/Sub_2_output_0, %onnx::Unsqueeze_186), scope: __main__.UNet::/__main__.Up::up4\n", " %onnx::Unsqueeze_188 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_6_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze_6\"](%/up4/Cast_3_output_0, %onnx::Unsqueeze_188), scope: __main__.UNet::/__main__.Up::up4\n", " %onnx::Unsqueeze_190 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}]()\n", " %/up4/Unsqueeze_7_output_0 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[onnx_name=\"/up4/Unsqueeze_7\"](%/up4/Sub_3_output_0, %onnx::Unsqueeze_190), scope: __main__.UNet::/__main__.Up::up4\n", " %/up4/Concat_1_output_0 : Long(4, strides=[1], device=cpu) = onnx::Concat[axis=0, onnx_name=\"/up4/Concat_1\"](%/up4/Unsqueeze_4_output_0, %/up4/Unsqueeze_5_output_0, %/up4/Unsqueeze_6_output_0, %/up4/Unsqueeze_7_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %onnx::Pad_193 : NoneType = prim::Constant(), scope: __main__.UNet::/__main__.Up::up4\n", " %/up4/Constant_6_output_0 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}, onnx_name=\"/up4/Constant_6\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Shape_4_output_0 : Long(1, strides=[1], device=cpu) = onnx::Shape[onnx_name=\"/up4/Shape_4\"](%/up4/Concat_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Gather_4_output_0 : Long(1, strides=[1], device=cpu) = onnx::Gather[axis=0, onnx_name=\"/up4/Gather_4\"](%/up4/Shape_4_output_0, %/up4/Constant_6_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_7_output_0 : Long(requires_grad=0, device=cpu) = onnx::Constant[value={8}, onnx_name=\"/up4/Constant_7\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Sub_4_output_0 : Long(1, strides=[1], device=cpu) = onnx::Sub[onnx_name=\"/up4/Sub_4\"](%/up4/Constant_7_output_0, %/up4/Gather_4_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Cast_4_output_0 : Long(4, strides=[1], device=cpu) = onnx::Cast[to=7, onnx_name=\"/up4/Cast_4\"](%/up4/Concat_1_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/ConstantOfShape_output_0 : Long(4, strides=[1], device=cpu) = onnx::ConstantOfShape[value={0}, onnx_name=\"/up4/ConstantOfShape\"](%/up4/Sub_4_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Concat_2_output_0 : Long(8, strides=[1], device=cpu) = onnx::Concat[axis=0, onnx_name=\"/up4/Concat_2\"](%/up4/Cast_4_output_0, %/up4/ConstantOfShape_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_8_output_0 : Long(2, strides=[1], device=cpu) = onnx::Constant[value=-1 2 [ CPULongType{2} ], onnx_name=\"/up4/Constant_8\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Reshape_output_0 : Long(4, 2, strides=[2, 1], device=cpu) = onnx::Reshape[allowzero=0, onnx_name=\"/up4/Reshape\"](%/up4/Concat_2_output_0, %/up4/Constant_8_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_9_output_0 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={0}, onnx_name=\"/up4/Constant_9\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_10_output_0 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={-1}, onnx_name=\"/up4/Constant_10\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_11_output_0 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={-9223372036854775807}, onnx_name=\"/up4/Constant_11\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_12_output_0 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={-1}, onnx_name=\"/up4/Constant_12\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Slice_output_0 : Long(4, 2, strides=[2, 1], device=cpu) = onnx::Slice[onnx_name=\"/up4/Slice\"](%/up4/Reshape_output_0, %/up4/Constant_10_output_0, %/up4/Constant_11_output_0, %/up4/Constant_9_output_0, %/up4/Constant_12_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Transpose_output_0 : Long(2, 4, strides=[4, 1], device=cpu) = onnx::Transpose[perm=[1, 0], onnx_name=\"/up4/Transpose\"](%/up4/Slice_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Constant_13_output_0 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={-1}, onnx_name=\"/up4/Constant_13\"](), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Reshape_1_output_0 : Long(8, strides=[1], device=cpu) = onnx::Reshape[allowzero=0, onnx_name=\"/up4/Reshape_1\"](%/up4/Transpose_output_0, %/up4/Constant_13_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Cast_5_output_0 : Long(8, strides=[1], device=cpu) = onnx::Cast[to=7, onnx_name=\"/up4/Cast_5\"](%/up4/Reshape_1_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Pad_output_0 : Float(*, *, *, *, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Pad[mode=\"constant\", onnx_name=\"/up4/Pad\"](%/up4/up/ConvTranspose_output_0, %/up4/Cast_5_output_0, %onnx::Pad_193), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:39:0\n", " %/up4/Concat_3_output_0 : Float(*, *, 256, 256, strides=[8388608, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Concat[axis=1, onnx_name=\"/up4/Concat_3\"](%/inc/conv/conv.5/Relu_output_0, %/up4/Pad_output_0), scope: __main__.UNet::/__main__.Up::up4 # /tmp/ipykernel_19565/1651252862.py:40:0\n", " %/up4/conv/conv/conv.0/Conv_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], onnx_name=\"/up4/conv/conv/conv.0/Conv\"](%/up4/Concat_3_output_0, %onnx::Conv_238, %onnx::Conv_239), scope: __main__.UNet::/__main__.Up::up4/__main__.DoubleConv::conv/torch.nn.modules.container.Sequential::conv/torch.nn.modules.conv.Conv2d::conv.0 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0\n", " %/up4/conv/conv/conv.2/Relu_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Relu[onnx_name=\"/up4/conv/conv/conv.2/Relu\"](%/up4/conv/conv/conv.0/Conv_output_0), scope: __main__.UNet::/__main__.Up::up4/__main__.DoubleConv::conv/torch.nn.modules.container.Sequential::conv/torch.nn.modules.activation.ReLU::conv.2 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/functional.py:1453:0\n", " %/up4/conv/conv/conv.3/Conv_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], onnx_name=\"/up4/conv/conv/conv.3/Conv\"](%/up4/conv/conv/conv.2/Relu_output_0, %onnx::Conv_241, %onnx::Conv_242), scope: __main__.UNet::/__main__.Up::up4/__main__.DoubleConv::conv/torch.nn.modules.container.Sequential::conv/torch.nn.modules.conv.Conv2d::conv.3 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0\n", " %/up4/conv/conv/conv.5/Relu_output_0 : Float(*, 64, 256, 256, strides=[4194304, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Relu[onnx_name=\"/up4/conv/conv/conv.5/Relu\"](%/up4/conv/conv/conv.3/Conv_output_0), scope: __main__.UNet::/__main__.Up::up4/__main__.DoubleConv::conv/torch.nn.modules.container.Sequential::conv/torch.nn.modules.activation.ReLU::conv.5 # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/functional.py:1453:0\n", " %/outc/conv/Conv_output_0 : Float(*, 1, 256, 256, strides=[65536, 65536, 256, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name=\"/outc/conv/Conv\"](%/up4/conv/conv/conv.5/Relu_output_0, %outc.conv.weight, %outc.conv.bias), scope: __main__.UNet::/__main__.OutConv::outc/torch.nn.modules.conv.Conv2d::conv # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0\n", " %output : Float(*, 1, 256, 256, strides=[65536, 65536, 256, 1], requires_grad=1, device=cpu) = onnx::Sigmoid[onnx_name=\"/outc/sigmoid/Sigmoid\"](%/outc/conv/Conv_output_0), scope: __main__.UNet::/__main__.OutConv::outc/torch.nn.modules.activation.Sigmoid::sigmoid # /home/gautham/.local/lib/python3.10/site-packages/torch/nn/modules/activation.py:294:0\n", " return (%output)\n", "\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/gautham/.local/lib/python3.10/site-packages/torch/onnx/utils.py:687: UserWarning: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. (Triggered internally at ../torch/csrc/jit/passes/onnx/shape_type_inference.cpp:1884.)\n", " _C._jit_pass_onnx_graph_shape_type_inference(\n", "/home/gautham/.local/lib/python3.10/site-packages/torch/onnx/utils.py:1178: UserWarning: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. (Triggered internally at ../torch/csrc/jit/passes/onnx/shape_type_inference.cpp:1884.)\n", " _C._jit_pass_onnx_graph_shape_type_inference(\n" ] } ], "source": [ "onnx_path = '../weights/model.onnx'\n", "\n", "torch.onnx.export(model,\n", " dummy_input,\n", " onnx_path,\n", " verbose=True,\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": [ "## Verifying the ONNX model" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import onnx\n", "\n", "onnx_model = onnx.load(onnx_path)\n", "onnx.checker.check_model(onnx_model)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Comparing ONNX Runtime and PyTorch results" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Exported model has been tested with ONNXRuntime, and the result looks good!\n" ] } ], "source": [ "import onnxruntime\n", "import numpy as np\n", "\n", "ort_session = onnxruntime.InferenceSession(onnx_path)\n", "\n", "def to_numpy(tensor):\n", " return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n", "\n", "ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(dummy_input)}\n", "ort_outs = ort_session.run(None, ort_inputs)\n", "\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!\")\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3.10.6 64-bit", "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.6" }, "orig_nbformat": 4, "vscode": { "interpreter": { "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1" } } }, "nbformat": 4, "nbformat_minor": 2 }