initial commit

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+flagged
+weights/water_bodies_model.pth
+venv

app.py

@@ -0,0 +1,11 @@
+import gradio as gr
+from functions import *
+
+title = "Water Body Segmentation - Image Segmentation PyTorch"
+examples = ['samples/image1.png', 'samples/image2.png', 'samples/image3.png', 'samples/image4.png', 'samples/image5.png']
+
+interface = gr.Interface(fn=predict, inputs=gr.Image(type= 'numpy').style(height= 256),
+            outputs= gr.Image(type = "numpy").style(height= 256),
+            examples= examples, title= title, css= '.gr-box {background-color: rgb(230 230 230);}')
+
+interface.launch()

functions.py

@@ -0,0 +1,36 @@
+import cv2
+import numpy as np
+import onnxruntime as rt
+
+def resize_preserve_aspect_ratio(image, size):
+    h, w = image.shape[:2]
+    if h > w:
+        image = cv2.resize(image, (size * w // h, size))
+    else:
+        image = cv2.resize(image, (size, size * h // w))
+    return image
+
+def predict(inp_image):
+
+    model_path = 'weights/model.onnx'
+    inp_dim = inp_image.shape[:2]
+
+    image = cv2.resize(inp_image, (256, 256))
+    image = np.array(image, dtype=np.float32) / 255.0
+    image = np.transpose(image, (2, 0, 1))
+    image = np.expand_dims(image, axis=0)
+
+    session = rt.InferenceSession(model_path)
+    input_name = session.get_inputs()[0].name
+    output_name = session.get_outputs()[0].name
+
+    pred_onx = session.run([output_name], {input_name: image.astype(np.float32)})[0]
+    pred_onx = pred_onx > 0.5
+    pred_onx = pred_onx * 255
+
+    pred_onx = cv2.resize(pred_onx[0, 0].astype(np.uint8) , (inp_dim[1], inp_dim[0]))
+    pred_onx = np.expand_dims(pred_onx, axis=2)
+    pred_onx = np.concatenate((pred_onx, pred_onx, pred_onx), axis=2)
+
+    output = resize_preserve_aspect_ratio(pred_onx, 256)
+    return output

notebooks/torch-to-onnx.ipynb

@@ -0,0 +1,389 @@
+{
+ "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.<locals>.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.<locals>.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
+}

requirements.txt

@@ -0,0 +1,67 @@
+aiohttp==3.8.3
+aiosignal==1.3.1
+anyio==3.6.2
+async-timeout==4.0.2
+attrs==22.1.0
+bcrypt==4.0.1
+certifi==2022.9.24
+cffi==1.15.1
+charset-normalizer==2.1.1
+click==8.1.3
+coloredlogs==15.0.1
+contourpy==1.0.6
+cryptography==38.0.3
+cycler==0.11.0
+fastapi==0.87.0
+ffmpy==0.3.0
+flatbuffers==22.10.26
+fonttools==4.38.0
+frozenlist==1.3.3
+fsspec==2022.11.0
+gradio==3.10.0
+h11==0.12.0
+httpcore==0.15.0
+httpx==0.23.0
+humanfriendly==10.0
+idna==3.4
+Jinja2==3.1.2
+kiwisolver==1.4.4
+linkify-it-py==1.0.3
+markdown-it-py==2.1.0
+MarkupSafe==2.1.1
+matplotlib==3.6.2
+mdit-py-plugins==0.3.1
+mdurl==0.1.2
+mpmath==1.2.1
+multidict==6.0.2
+numpy==1.23.4
+onnxruntime==1.13.1
+opencv-python==4.6.0.66
+orjson==3.8.1
+packaging==21.3
+pandas==1.5.1
+paramiko==2.12.0
+Pillow==9.3.0
+protobuf==4.21.9
+pycparser==2.21
+pycryptodome==3.15.0
+pydantic==1.10.2
+pydub==0.25.1
+PyNaCl==1.5.0
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytz==2022.6
+PyYAML==6.0
+requests==2.28.1
+rfc3986==1.5.0
+six==1.16.0
+sniffio==1.3.0
+starlette==0.21.0
+sympy==1.11.1
+typing_extensions==4.4.0
+uc-micro-py==1.0.1
+urllib3==1.26.12
+uvicorn==0.19.0
+websockets==10.4
+yarl==1.8.1

weights/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7e5833c8c9f26b18a074c313c11c4db5e569cdbf65357ae057caaa895e2bfbb
+size 1837091