Segmentation_Training_Script_1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fresh-jesus",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim \n",
    "import torchvision\n",
    "from torchvision import models\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torch.utils.data as utils\n",
    "from torchvision import transforms\n",
    "import torch.nn.functional as F\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "from Data_Retriever import ImageData\n",
    "from Resnet_UNet import ResNetUNet\n",
    "from Metrics import Metrics\n",
    "import time\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "digital-japan",
   "metadata": {},
   "outputs": [],
   "source": [
    "CUDA_LAUNCH_BLOCKING=\"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "unlimited-absence",
   "metadata": {},
   "outputs": [],
   "source": [
    "tr = pd.read_csv('train.csv')\n",
    "tr = tr[tr['ClassId'] == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "integral-france",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df, test_df = train_test_split(tr, test_size=0.33, random_state=123)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "tough-session",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "600"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "technological-remove",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transforms = transforms.Compose([\n",
    "                                transforms.Resize((256, 256)),\n",
    "                                transforms.ToTensor()])\n",
    "                                #transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                                     #[0.229, 0.224, 0.225])])\n",
    "train_data = ImageData(df = train_df, transform = train_transforms, subset=\"train\")\n",
    "train_loader = DataLoader(dataset = train_data, batch_size=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "common-mailing",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_transforms = transforms.Compose([\n",
    "                                transforms.Resize((256, 256)),\n",
    "                                transforms.ToTensor()])\n",
    "                                #transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                                     #[0.229, 0.224, 0.225])])\n",
    "test_data = ImageData(df = test_df, transform = test_transforms, subset=\"test\")\n",
    "test_loader = DataLoader(dataset = train_data, batch_size=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "raising-elimination",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "model = ResNetUNet(n_class=1)\n",
    "criterion = nn.BCEWithLogitsLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), weight_decay=1e-4, lr = 0.001, momentum=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "listed-italy",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ResNetUNet(\n",
      "  (base_model): ResNet(\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",
      "    (relu): ReLU(inplace)\n",
      "    (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "    (layer1): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer2): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer3): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (downsample): Sequential(\n",
      "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (layer4): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\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): BasicBlock(\n",
      "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "    (fc): Linear(in_features=512, out_features=1000, bias=True)\n",
      "  )\n",
      "  (layer0): Sequential(\n",
      "    (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
      "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU(inplace)\n",
      "  )\n",
      "  (layer0_1x1): Sequential(\n",
      "    (0): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (layer1): Sequential(\n",
      "    (0): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
      "    (1): Sequential(\n",
      "      (0): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "      (1): BasicBlock(\n",
      "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (relu): ReLU(inplace)\n",
      "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (layer1_1x1): Sequential(\n",
      "    (0): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer2_1x1): Sequential(\n",
      "    (0): Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer3_1x1): Sequential(\n",
      "    (0): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (layer4): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\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): BasicBlock(\n",
      "      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer4_1x1): Sequential(\n",
      "    (0): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (upsample): Upsample(scale_factor=2.0, mode=bilinear)\n",
      "  (conv_up3): Sequential(\n",
      "    (0): Conv2d(768, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (conv_up2): Sequential(\n",
      "    (0): Conv2d(640, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (conv_up1): Sequential(\n",
      "    (0): Conv2d(320, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (conv_up0): Sequential(\n",
      "    (0): Conv2d(320, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (conv_original_size0): Sequential(\n",
      "    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (conv_original_size1): Sequential(\n",
      "    (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (conv_original_size2): Sequential(\n",
      "    (0): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU(inplace)\n",
      "  )\n",
      "  (conv_last): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "improved-patrol",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "DEVICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "084c36db",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = Metrics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "218664a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\SIDDHARTA\\anaconda3\\envs\\pytorchenv\\lib\\site-packages\\torch\\nn\\functional.py:1386: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
      "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model training\n",
      "epoch:  0 loss: 0.00141796     dice_loss: 0.99  dice_coeff  8.57e-05\n",
      "Model Validation\n",
      "epoch:  0 loss: 0.00141300         dice_loss: 0.99  dice_coeff  8.57e-05\n",
      "Model training\n",
      "epoch:  1 loss: 0.00096414     dice_loss: 0.99  dice_coeff  8.71e-05\n",
      "Model Validation\n",
      "epoch:  1 loss: 0.00096329         dice_loss: 0.99  dice_coeff  8.71e-05\n",
      "Model training\n",
      "epoch:  2 loss: 0.00084090     dice_loss: 0.98  dice_coeff  8.63e-05\n",
      "Model Validation\n",
      "epoch:  2 loss: 0.00084051         dice_loss: 0.98  dice_coeff  8.63e-05\n",
      "Model training\n",
      "epoch:  3 loss: 0.00079373     dice_loss: 0.98  dice_coeff  8.61e-05\n",
      "Model Validation\n",
      "epoch:  3 loss: 0.00079348         dice_loss: 0.98  dice_coeff  8.61e-05\n",
      "Model training\n",
      "epoch:  4 loss: 0.00076791     dice_loss: 0.97  dice_coeff  8.59e-05\n",
      "Model Validation\n",
      "epoch:  4 loss: 0.00076775         dice_loss: 0.97  dice_coeff  8.59e-05\n",
      "\n",
      "Duration: 7591 seconds\n",
      "Wall time: 2h 6min 31s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "EPOCHS = 5\n",
    "epochs = 0\n",
    "start_time = time.time()\n",
    "\n",
    "#configure the model for GPU training\n",
    "model = model.to(DEVICE)\n",
    "train_losses = []\n",
    "test_losses = []\n",
    "train_dice_loss = []\n",
    "test_dice_loss = []\n",
    "train_dice_coeff = []\n",
    "test_dice_coeff = []\n",
    "\n",
    "PATH = 'segmentation_1_checkpoint.pth'\n",
    "\n",
    "if os.path.exists(PATH):\n",
    "    checkpoint = torch.load(PATH)\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "    epochs = checkpoint['epoch']\n",
    "    train_losses = checkpoint['train_loss']\n",
    "    test_losses = checkpoint['test_loss']\n",
    "    train_dice_loss = checkpoint['train_dice_loss']\n",
    "    test_dice_loss = checkpoint['test_dice_loss']\n",
    "    train_dice_coeff = checkpoint['train_dice_coeff']\n",
    "    test_dice_coeff = checkpoint['test_dice_coeff']\n",
    "    \n",
    "    \n",
    "for i in range(epochs, EPOCHS):     \n",
    "    trn_corr = 0\n",
    "    tst_corr = 0\n",
    "    model.train()         \n",
    "    for b, (data, target) in enumerate(train_loader): \n",
    "        b += 1\n",
    "        data, target = data.cuda(), target.cuda()\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)  \n",
    "        loss = criterion(output, target)\n",
    "        dice_loss = metrics.diceloss(output, target)\n",
    "        dice_coeff = metrics.dice_coeff(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "    print(f'Model training')\n",
    "    print(f'epoch: {i:2} loss: {loss.item():5.8f} \\\n",
    "    dice_loss:{dice_loss.item():5.2f}  dice_coeff{dice_coeff.item():10.3}')  \n",
    "            \n",
    "\n",
    "    train_losses.append(loss)\n",
    "    train_dice_loss.append(dice_loss)\n",
    "    train_dice_coeff.append(dice_coeff)\n",
    "    \n",
    "     # Run the testing batches\n",
    "    with torch.no_grad():\n",
    "        for b, (X_test, y_test) in enumerate(test_loader):\n",
    "            b+=1\n",
    "            \n",
    "            X_test = X_test.to(DEVICE)\n",
    "            y_test = y_test.to(DEVICE)\n",
    "            # Apply the model\n",
    "            y_val = model(X_test)\n",
    "\n",
    "\n",
    "        loss = criterion(y_val, y_test)\n",
    "        dice_loss = metrics.diceloss(y_val, y_test)\n",
    "        dice_coeff = metrics.dice_coeff(y_val, y_test)\n",
    "        test_losses.append(loss)\n",
    "        test_dice_loss.append(dice_loss)\n",
    "        test_dice_coeff.append(dice_coeff)\n",
    "\n",
    "        print(f'Model Validation')\n",
    "        print(f'epoch: {i:2} loss: {loss.item():5.8f} \\\n",
    "        dice_loss:{dice_loss.item():5.2f}  dice_coeff{dice_coeff.item():10.3}')  \n",
    "    \n",
    "    checkpoint = {'epoch': i,\n",
    "                  'state_dict': model.state_dict(),\n",
    "                  'optimizer' : optimizer.state_dict(),\n",
    "                  'train_loss' : train_losses,\n",
    "                  'test_loss' : test_losses,\n",
    "                  'train_dice_loss' : train_dice_loss,\n",
    "                  'test_dice_loss' : test_dice_loss,\n",
    "                  'train_dice_coeff':train_dice_coeff,\n",
    "                  'test_dice_coeff':test_dice_coeff}\n",
    "\n",
    "    torch.save(checkpoint, PATH)\n",
    "\n",
    "print(f'\\nDuration: {time.time() - start_time:.0f} seconds') # print the time elapsed\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "mature-knight",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'Segmentation_1.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "residential-advertiser",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.to(DEVICE)\n",
    "model.load_state_dict(torch.load('Segmentation_1.pt'));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d84c3ec7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}