Inference Notebook with TRT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### UNet Inference kernel\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "import pdb\n",
    "import os\n",
    "import cv2\n",
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import torch.backends.cudnn as cudnn\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from albumentations import (Normalize, Compose)\n",
    "from albumentations.pytorch import ToTensor\n",
    "import torch.utils.data as data\n",
    "import tensorrt as trt\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestDataset(Dataset):\n",
    "    '''Dataset for test prediction'''\n",
    "    def __init__(self, root, df, mean, std):\n",
    "        self.root = root\n",
    "        #df['ImageId'] = df['ImageId_ClassId'].apply(lambda x: x.split('_')[0])\n",
    "        self.fnames = df['ImageId'].unique().tolist()\n",
    "        self.num_samples = len(self.fnames)\n",
    "        self.transform = Compose(\n",
    "            [\n",
    "                Normalize(mean=mean, std=std, p=1),\n",
    "                ToTensor(),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        fname = self.fnames[idx]\n",
    "        path = os.path.join(self.root, fname)\n",
    "        image = cv2.imread(path)\n",
    "        images = self.transform(image=image)[\"image\"]\n",
    "        return fname, images\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best_threshold 0.5\n"
     ]
    }
   ],
   "source": [
    "# initialize test dataloader\n",
    "sample_submission_path = './severstal-steel-defect-detection/sample_submission.csv'\n",
    "test_data_folder = \"./severstal-steel-defect-detection/test_images\"\n",
    "best_threshold = 0.5\n",
    "num_workers = 2\n",
    "batch_size = 4\n",
    "print('best_threshold', best_threshold)\n",
    "min_size = 3500\n",
    "mean = (0.485, 0.456, 0.406)\n",
    "std = (0.229, 0.224, 0.225)\n",
    "df = pd.read_csv(sample_submission_path)\n",
    "testset = DataLoader(\n",
    "    TestDataset(test_data_folder, df, mean, std),\n",
    "    batch_size=batch_size,\n",
    "    shuffle=False,\n",
    "    num_workers=num_workers,\n",
    "    pin_memory=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load PyT Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/torch/serialization.py:593: SourceChangeWarning: source code of class 'segmentation_models_pytorch.encoders.senet.SENetEncoder' has changed. you can retrieve the original source code by accessing the object's source attribute or set `torch.nn.Module.dump_patches = True` and use the patch tool to revert the changes.\n",
      "  warnings.warn(msg, SourceChangeWarning)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Unet(\n",
       "  (encoder): SENetEncoder(\n",
       "    (layer0): Sequential(\n",
       "      (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu1): ReLU(inplace=True)\n",
       "      (pool): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=True)\n",
       "    )\n",
       "    (layer1): Sequential(\n",
       "      (0): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (2): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(256, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer2): Sequential(\n",
       "      (0): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (2): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (3): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(512, 32, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(32, 512, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer3): Sequential(\n",
       "      (0): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (2): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (3): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (4): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (5): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (6): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (7): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (8): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (9): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (10): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (11): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (12): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (13): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (14): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (15): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (16): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (17): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (18): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (19): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (20): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (21): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (22): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(1024, 64, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(64, 1024, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer4): Sequential(\n",
       "      (0): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(1024, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(2048, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "      (2): SEResNeXtBottleneck(\n",
       "        (conv1): Conv2d(2048, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "        (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (se_module): SEModule(\n",
       "          (avg_pool): AdaptiveAvgPool2d(output_size=1)\n",
       "          (fc1): Conv2d(2048, 128, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (relu): ReLU(inplace=True)\n",
       "          (fc2): Conv2d(128, 2048, kernel_size=(1, 1), stride=(1, 1))\n",
       "          (sigmoid): Sigmoid()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (decoder): UnetDecoder(\n",
       "    (center): Identity()\n",
       "    (blocks): ModuleList(\n",
       "      (0): DecoderBlock(\n",
       "        (conv1): Conv2dReLU(\n",
       "          (0): Conv2d(3072, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention1): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "        (conv2): Conv2dReLU(\n",
       "          (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention2): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "      )\n",
       "      (1): DecoderBlock(\n",
       "        (conv1): Conv2dReLU(\n",
       "          (0): Conv2d(768, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention1): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "        (conv2): Conv2dReLU(\n",
       "          (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention2): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "      )\n",
       "      (2): DecoderBlock(\n",
       "        (conv1): Conv2dReLU(\n",
       "          (0): Conv2d(384, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention1): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "        (conv2): Conv2dReLU(\n",
       "          (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention2): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "      )\n",
       "      (3): DecoderBlock(\n",
       "        (conv1): Conv2dReLU(\n",
       "          (0): Conv2d(128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention1): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "        (conv2): Conv2dReLU(\n",
       "          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention2): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "      )\n",
       "      (4): DecoderBlock(\n",
       "        (conv1): Conv2dReLU(\n",
       "          (0): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention1): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "        (conv2): Conv2dReLU(\n",
       "          (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): ReLU(inplace=True)\n",
       "        )\n",
       "        (attention2): Attention(\n",
       "          (attention): Identity()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (segmentation_head): SegmentationHead(\n",
       "    (0): Conv2d(16, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (1): Identity()\n",
       "    (2): Activation(\n",
       "      (activation): Identity()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize mode and load trained weights\n",
    "ckpt_path = \"./models/modelmodel.pth\"\n",
    "device = torch.device(\"cuda\")\n",
    "state = torch.load(ckpt_path)\n",
    "model_pyt = torch.load(ckpt_path)\n",
    "model_pyt.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference on PyT model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1377/1377 [02:28<00:00,  9.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time for predictions:  148.0280213356018 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# start prediction\n",
    "import time\n",
    "start = time.time()\n",
    "predictions = []\n",
    "for i, batch in enumerate(tqdm(testset)):\n",
    "    fnames, images = batch\n",
    "    batch_preds = torch.sigmoid(model_pyt(images.to(device)))\n",
    "    batch_preds = batch_preds.detach().cpu().numpy()\n",
    "end = time.time()\n",
    "inference_pyt=end-start\n",
    "print(\"Time for predictions: \",inference_pyt,'s')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load TRT model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import common2\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.WARNING)\n",
    "trt_runtime = trt.Runtime(TRT_LOGGER)\n",
    "engine_path = \"my_engine.trt\"\n",
    "model = common2.load_engine(trt_runtime,engine_path)\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "context = model.create_execution_context()\n",
    "\n",
    "inputs, outputs, bindings, stream = common2.allocate_buffers(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import listdir\n",
    "from os.path import isfile, join\n",
    "import numpy\n",
    "import cv2\n",
    "\n",
    "mypath='./severstal-steel-defect-detection/test_images/'\n",
    "batch_size=1377\n",
    "onlyfiles = [ f for f in listdir(mypath) if isfile(join(mypath,f)) ]\n",
    "image = numpy.empty(batch_size, dtype=object)\n",
    "for n in range(0, batch_size):\n",
    "    image[n] = cv2.imread( join(mypath,onlyfiles[n]) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference on TRT model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inference time:  3.0406620502471924 s\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "for n in range(0, batch_size):\n",
    "    inputs[0].host = image[n]\n",
    "    trt_outputs = common2.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)\n",
    "end = time.time()\n",
    "inference_trt = end-start\n",
    "print(\"Inference time: \",inference_trt, \"s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Total SpeedUp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Speedup = 48.6828259403467 x\n"
     ]
    }
   ],
   "source": [
    "print(\"Speedup =\",inference_pyt/inference_trt,\"x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Refrences:\n",
    "\n",
    "Few kernels from which I've borrowed some code:\n",
    "\n",
    "* https://www.kaggle.com/amanooo/defect-detection-starter-u-net\n",
    "* https://www.kaggle.com/go1dfish/clear-mask-visualization-and-simple-eda"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}