second_ty.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/Arshad221b/MedSeg/blob/main/second_ty.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "4CHv0U1wPCG0"
   },
   "outputs": [],
   "source": [
    "import random \n",
    "import torch \n",
    "from torch.utils import data \n",
    "import torchvision.transforms.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision.transforms import transforms\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rud7WqEHb8Ve",
    "outputId": "dd5d2c59-f066-4924-db66-f89a0b2af607"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: nibabel in /home/arshad/Downloads/venv/lib/python3.11/site-packages (5.1.0)\n",
      "Requirement already satisfied: numpy>=1.19 in /home/arshad/Downloads/venv/lib/python3.11/site-packages (from nibabel) (1.24.2)\n",
      "Requirement already satisfied: packaging>=17 in /home/arshad/Downloads/venv/lib/python3.11/site-packages (from nibabel) (23.1)\n"
     ]
    }
   ],
   "source": [
    "! pip install nibabel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "9Q0x1-i3cEU4"
   },
   "outputs": [],
   "source": [
    "import nibabel as nib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "mrdqqcEFPNP8"
   },
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cuda'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B7727guIDvSg"
   },
   "source": [
    "# Data Collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "ioZcoim2PRdO"
   },
   "outputs": [],
   "source": [
    "# ! unzip /content/drive/MyDrive/AMOS/amos22.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sw0PdPCnDyZJ"
   },
   "source": [
    "# Data Loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "U9DyYcKammlA"
   },
   "outputs": [],
   "source": [
    "IMAGE_SIZE = 32\n",
    "BATCH_SIZE = 2\n",
    "NUM_CLASS = 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "KRWGRtEvZbZK"
   },
   "outputs": [],
   "source": [
    "import cv2\n",
    "from glob import glob\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Q_7YVKIA-InO"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "UoQ6S3nUPGJG"
   },
   "outputs": [],
   "source": [
    "class AmosDataLoader(data.Dataset):\n",
    "  def __init__(\n",
    "      self, \n",
    "      input_paths: list, \n",
    "      target_paths: list, \n",
    "      transform_input = None, \n",
    "      transform_target = None\n",
    "  ): \n",
    "\n",
    "    self.input_paths      = input_paths\n",
    "    self.target_paths     = target_paths\n",
    "    self.transform_input  = transform_input\n",
    "    self.transform_target = transform_target\n",
    "\n",
    "  def __len__(self):\n",
    "    return len(self.input_paths)\n",
    "\n",
    "  def preprocess_img_input(self, input_im):\n",
    "    # z_factor_input      = input_im.shape[2]*int(input_im.shape[0]/IMAGE_SIZE)**2\n",
    "    input_im            = np.stack((input_im,)*3, axis=-1)\n",
    "    input_im            = torch.tensor(input_im).float()/255\n",
    "    # print(input_im.shape)\n",
    "    input_im            = input_im.permute(3,2,0,1)\n",
    "    \n",
    "    input_im            = input_im.unsqueeze(0)\n",
    "    # print('input_shape before inter',input_im.shape)\n",
    "    output_size_input   = (82, IMAGE_SIZE, IMAGE_SIZE)\n",
    "    input_im            = F.interpolate(input_im, size=output_size_input, mode='trilinear', align_corners=False)\n",
    "    input_im            = input_im#.squeeze(0)\n",
    "    return input_im\n",
    "\n",
    "  def preprocess_output(self, output_im):\n",
    "    # z_factor_output      = output_im.shape[2]*int(output_im.shape[0]/IMAGE_SIZE)**2\n",
    "    # output_im            = torch.tensor(output_im).float()/255\n",
    "    # print('output shape before mask', output_im.shape)\n",
    "    mask_cat              = np.zeros((NUM_CLASS, *output_im.shape), dtype=np.float32)\n",
    "    for i in range(NUM_CLASS):\n",
    "        mask_cat[i][output_im == i] = 1\n",
    "    # output_im = torch.tensor(mask_cat).float()\n",
    "    output_im             = torch.tensor(mask_cat).float()/255\n",
    "    # print(output_im.shape)\n",
    "    output_im             = output_im.permute(0,2,3,1)\n",
    "    # print(output_im.shape)\n",
    "    output_im             = output_im.unsqueeze(0)#.unsqueeze(0)\n",
    "    # print('output shape befor inter',output_im.shape)\n",
    "    output_size_input     = (82, IMAGE_SIZE, IMAGE_SIZE)\n",
    "    output_im             = F.interpolate(output_im, size=output_size_input, mode='trilinear', align_corners=False)\n",
    "    output_im             = output_im#.squeeze(0)\n",
    "    \n",
    "    \n",
    "    # print('output shape final', output_im.shape)\n",
    "    return output_im\n",
    "\n",
    "  def __getitem__(self,x):\n",
    "    input_image   = self.input_paths[x]\n",
    "    target_image  = self.target_paths[x]\n",
    "    input_im      = nib.load(input_image).get_fdata()\n",
    "    target_im     = nib.load(target_image).get_fdata()\n",
    "\n",
    "    return input_im, target_im\n",
    "    \n",
    "  def collate_fn(self, batch):\n",
    "    # print(len(batch[0][0]))\n",
    "    im_ins, im_outs = [], []\n",
    "    for im_in, im_out  in batch: \n",
    "      im_in = self.preprocess_img_input(im_in)\n",
    "      im_out = self.preprocess_output(im_out)\n",
    "      # print(im_in.shape, im_out.shape)\n",
    "      im_ins.append(im_in)\n",
    "      im_outs.append(im_out)\n",
    "\n",
    "    # print(torch.tensor(im_ins).shape)\n",
    "    return torch.cat(im_ins, dim = 0).to(device), torch.cat(im_outs, dim= 0).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "hrVS5thiZtF_"
   },
   "outputs": [],
   "source": [
    "path = '/home/arshad/Downloads/amos22/amos22'\n",
    "\n",
    "\n",
    "input_paths   = sorted(glob(os.path.join(path, \"imagesVa\",\"*.nii.gz\")))\n",
    "target_paths  = sorted(glob(os.path.join(path, \"labelsVa\",\"*.nii.gz\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_paths = input_paths[:25]\n",
    "target_paths = target_paths[:25]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['labelsTs',\n",
       " 'imagesTr',\n",
       " 'imagesVa',\n",
       " 'labelsVa',\n",
       " 'dataset.json',\n",
       " 'readme.md',\n",
       " '.DS_Store',\n",
       " 'imagesTs',\n",
       " 'labelsTr']"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.listdir(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "id": "fioLA0sqaUru"
   },
   "outputs": [],
   "source": [
    "train_dl      = AmosDataLoader(input_paths, target_paths)\n",
    "train_loader  = DataLoader(train_dl, batch_size = BATCH_SIZE, drop_last= True, collate_fn=train_dl.collate_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "M1Lu0kq5b7b0",
    "outputId": "fb203998-7787-4815-d8f0-e8199b5cc2ce"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(768, 768, 101) (768, 768, 101)\n"
     ]
    }
   ],
   "source": [
    "for d,k in train_dl: \n",
    "  # # print(len(d))}\n",
    "  print(d.shape, k.shape)\n",
    "  break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "XtIHG6Xq20_2"
   },
   "outputs": [],
   "source": [
    "# def preprocess_output( output_im):\n",
    "#     z_factor_output      = output_im.shape[2]*int(output_im.shape[0]/IMAGE_SIZE)**2\n",
    "#     # output_im            = torch.tensor(output_im).float()/255\n",
    "#     mask_cat              = np.zeros((NUM_CLASS, *output_im.shape), dtype=np.float32)\n",
    "#     for i in range(NUM_CLASS):\n",
    "#         mask_cat[i][output_im == i] = 1\n",
    "\n",
    "#     output_im             = torch.tensor(mask_cat).float()/255\n",
    "#     output_im             = output_im.permute(3,2,0,1)\n",
    "#     # print(output_im.shape)\n",
    "#     output_im             = output_im.unsqueeze(0)\n",
    "#     output_size_input     = (15, IMAGE_SIZE, IMAGE_SIZE)\n",
    "#     output_im             = F.interpolate(output_im, size=output_size_input, mode='trilinear', align_corners=False)\n",
    "#     output_im             = output_im#.squeeze(0)\n",
    "#     return output_im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "id": "W5kKlwtntCQ0"
   },
   "outputs": [],
   "source": [
    "# import torchvision.transforms.functional as TF \n",
    "# target_image = target_paths[0]\n",
    "# mask_nii     = nib.load(target_image)\n",
    "# mask_data = mask_nii.get_fdata()\n",
    "\n",
    "# o = preprocess_output(mask_data)\n",
    "# # target_im  = torch.tensor(target_im)\n",
    "\n",
    "# num_classes = 15\n",
    "\n",
    "# # Convert the mask to categorical values\n",
    "# mask_cat = np.zeros((num_classes, *mask_data.shape), dtype=np.float32)\n",
    "# for i in range(num_classes):\n",
    "#     mask_cat[i][mask_data == i] = 1\n",
    "\n",
    "# # Save the categorical mask as a new NIfTI image\n",
    "# mask_cat_nii = nib.Nifti1Image(mask_cat, mask_nii.affine, mask_nii.header)\n",
    "# # nib.save(mask_cat_nii, \"mask_cat.nii.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "id": "OM2hpCtO3Ane"
   },
   "outputs": [],
   "source": [
    "# o.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "id": "THpaDPJT3Wmd"
   },
   "outputs": [],
   "source": [
    "o = 0\n",
    "mask_data = 0\n",
    "mask_nii = 0 \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "id": "VRpiVvhQvz88"
   },
   "outputs": [],
   "source": [
    "# mask_cat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 435
    },
    "id": "N1fd1WQmup8Y",
    "outputId": "2ad191c2-87fa-4213-e319-f39ba8e1fa26"
   },
   "outputs": [],
   "source": [
    "# import nibabel as nib\n",
    "\n",
    "# test_load = nib.load('/content/amos22/labelsVa/amos_0008.nii.gz').get_fdata()\n",
    "# test = test_load[:,:,21]\n",
    "# plt.imshow(test)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "id": "4nAwMVJXbPjg"
   },
   "outputs": [],
   "source": [
    "# for i, j in train_dl: \n",
    "#   print(i.shape, j.shape)\n",
    "#   i = i.permute(2,1,0)\n",
    "#   plt.imshow(i[:,:,59].unsqeeze(0).to('cpu'))\n",
    "#   plt.show()N\n",
    "#   # print(i)\n",
    "#   break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "id": "9hskYVprmMAw"
   },
   "outputs": [],
   "source": [
    "# import torch.nn.functional as F\n",
    "\n",
    "# # Assume tensor of shape (batch_size, channels, depth, height, width)\n",
    "# x = torch.randn(1, 1, 90, 768, 768)\n",
    "\n",
    "# # Specify the output size of the tensor after interpolation\n",
    "# output_size = (810,  256, 256)\n",
    "\n",
    "# # Interpolate the tensor to the desired output size using trilinear interpolation\n",
    "# y = F.interpolate(x, size=output_size, mode='trilinear', align_corners=False)\n",
    "\n",
    "# # The shape of the resulting tensor should be (batch_size, channels, depth, height, width)\n",
    "# print(y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iH43xw0BOoQe"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "id": "tZ_azzwzaifB"
   },
   "outputs": [],
   "source": [
    "# a,b = next(iter(train_loader))\n",
    "# a.shape, b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PPB3L7WsaoVa",
    "outputId": "f20aff99-d1f8-4b9d-bb54-ef5d1cdf5d1f"
   },
   "outputs": [],
   "source": [
    "# input_paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vGsVE4joDrAp"
   },
   "source": [
    "# Model Building\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "id": "tAXk92Ooa0YO"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class DoubleConv3D(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv3d(in_channels, out_channels, 3, padding=1),\n",
    "            nn.BatchNorm3d(out_channels),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv3d(out_channels, out_channels, 3, padding=1),\n",
    "            nn.BatchNorm3d(out_channels),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "class Down3D(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.mpconv = nn.Sequential(\n",
    "            nn.MaxPool3d(2),\n",
    "            DoubleConv3D(in_channels, out_channels)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mpconv(x)\n",
    "\n",
    "class Up3D(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, bilinear=True):\n",
    "        super().__init__()\n",
    "\n",
    "        if bilinear:\n",
    "            self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)\n",
    "        else:\n",
    "            self.up = nn.ConvTranspose3d(in_channels, in_channels//2, kernel_size=2, stride=2)\n",
    "\n",
    "        self.conv = DoubleConv3D(in_channels, out_channels)\n",
    "\n",
    "    def forward(self, x1, x2):\n",
    "        # print(x1.shape, x2.shape)\n",
    "        x1    = self.up(x1)\n",
    "        # print(x1.shape)\n",
    "        diffZ = x2.size()[2] - x1.size()[2]\n",
    "        diffY = x2.size()[3] - x1.size()[3]\n",
    "        diffX = x2.size()[4] - x1.size()[4]\n",
    "        x1    = F.pad(x1, (diffX // 2, diffX - diffX // 2,\n",
    "                        diffY // 2, diffY - diffY // 2,\n",
    "                        diffZ // 2, diffZ - diffZ // 2))\n",
    "        x     = torch.cat([x2, x1], dim=1)\n",
    "        return self.conv(x)\n",
    "\n",
    "class OutConv3D(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(OutConv3D, self).__init__()\n",
    "        self.conv = nn.Conv3d(in_channels, out_channels, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "class UNet3D(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, bilinear=False):\n",
    "        super().__init__()\n",
    "        self.in_channels  = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.bilinear     = bilinear\n",
    "\n",
    "        self.conv1    = DoubleConv3D(in_channels, 64)\n",
    "        self.down1    = Down3D(64, 128)\n",
    "        self.down2    = Down3D(128, 256)\n",
    "        self.down3    = Down3D(256, 512)\n",
    "        self.down4    = Down3D(512, 1024)\n",
    "        self.up1      = Up3D(1024, 512, bilinear)\n",
    "        self.up2      = Up3D(512, 256, bilinear)\n",
    "        self.up3      = Up3D(256, 128, bilinear)\n",
    "        self.up4      = Up3D(128, 64, bilinear)\n",
    "        self.outconv  = OutConv3D(64, out_channels)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # print(x.shape)\n",
    "        # x = x.unsqueeze(1)\n",
    "        x1 = self.conv1(x)\n",
    "        x2 = self.down1(x1)\n",
    "        x3 = self.down2(x2)\n",
    "        x4 = self.down3(x3)\n",
    "        x5 = self.down4(x4)\n",
    "        # print(x5.shape, x4.shape)\n",
    "        x6 = self.up1(x5, x4)\n",
    "        x7 = self.up2(x6, x3)\n",
    "        x8 = self.up3(x7, x2) \n",
    "        x9 = self.up4(x8, x1)\n",
    "        output= self.outconv(x9)\n",
    "        # print(x6.shape)\n",
    "        # up network\n",
    "\n",
    "        return output\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "id": "WYdfm81hGCRD"
   },
   "outputs": [],
   "source": [
    "model = UNet3D(3, 15).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DUNIpQ1h0NrO"
   },
   "source": [
    "# Model Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oX4l1JVM0NQh"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "id": "itTQ9SMSGHqd"
   },
   "outputs": [],
   "source": [
    "# im= torch.zeros((2, 3,82,128,128))\n",
    "# UNet3D(3,15).forward(im)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "id": "IDGOJGbSGT3f"
   },
   "outputs": [],
   "source": [
    "\n",
    "# !pip install torchsummary\n",
    "# from torchsummary import summary\n",
    "# summary(model, input_size=(3, 82, 128, 128), device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "id": "avXePzyVKpha"
   },
   "outputs": [],
   "source": [
    "\n",
    "def train_batch(data, model, optimizer, criteria):\n",
    "    model.train()\n",
    "    \n",
    "    ims_in, ims_out = data\n",
    "    # print(ims_in.shape)\n",
    "    optimizer.zero_grad()\n",
    "    # print(ims_in.shape)\n",
    "    pred_img = model(ims_in)    \n",
    "    # print(pred_img.shape)\n",
    "    # gender_criterion, age_criterion = criteria\n",
    "    # gender_loss = gender_criterion(pred_gender.squeeze(), gender)\n",
    "    # age_loss = age_criterion(pred_age.squeeze(), age)\n",
    "    total_loss = criteria(pred_img, ims_out)\n",
    "    total_loss.backward()\n",
    "    optimizer.step()\n",
    "    return total_loss\n",
    "\n",
    "# def validate_batch(data, model, criteria):\n",
    "#     model.eval()\n",
    "#     ims, age, gender = data\n",
    "#     with torch.no_grad():\n",
    "#       pred_gender, pred_age = model(ims)\n",
    "#     gender_criterion, age_criterion = criteria\n",
    "#     gender_loss = gender_criterion(pred_gender.squeeze(), gender)\n",
    "#     age_loss = age_criterion(pred_age.squeeze(), age)\n",
    "#     total_loss = gender_loss + age_loss\n",
    "#     pred_gender = (pred_gender > 0.5).squeeze()\n",
    "#     gender_acc = (pred_gender == gender).float().sum()\n",
    "#     age_mae = torch.abs(age - pred_age).float().sum()\n",
    "#     return total_loss, gender_acc, age_mae"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "id": "ekD7Nc3cSoar"
   },
   "outputs": [],
   "source": [
    "n_epoch = 10\n",
    "criteria = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "WoSuiFsNTs76",
    "outputId": "372dea96-5e63-4b1e-9bd2-949c827373fe"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "avg_loss tensor(0.0785, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "1\n",
      "avg_loss tensor(0.0271, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "2\n",
      "avg_loss tensor(0.0150, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "3\n",
      "avg_loss tensor(0.0091, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "4\n",
      "avg_loss tensor(0.0055, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "5\n",
      "avg_loss tensor(0.0033, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "6\n",
      "avg_loss tensor(0.0020, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "7\n",
      "avg_loss tensor(0.0013, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "8\n",
      "avg_loss tensor(0.0009, device='cuda:0', grad_fn=<DivBackward0>)\n",
      "9\n",
      "avg_loss tensor(0.0007, device='cuda:0', grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "train_loss = []\n",
    "for epoch in range(n_epoch):\n",
    "  print(epoch)\n",
    "  epoch_loss = []\n",
    "\n",
    "  for ix, ims in enumerate(train_loader):\n",
    "    loss = train_batch(ims, model, optimizer, criteria)\n",
    "    epoch_loss.append(loss)\n",
    "  print('avg_loss', sum(epoch_loss)/len(epoch_loss))\n",
    "  train_loss.append(sum(epoch_loss)/len(epoch_loss))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xtnNM1vRVEKx"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyPauRa/mx8GFCDYnEHvCNI0",
   "collapsed_sections": [
    "B7727guIDvSg"
   ],
   "include_colab_link": true,
   "mount_file_id": "15Fdb5Qh4O08i6gsGbw8ZCiMwIXiRtFH5",
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}