B7_SegFormer_mmseg.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Kiểm tra GPU được cấp phát</h3>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Cài đặt thư viện MMSegmentation</h3>\n",
    "</div>\n",
    "Tham khảo: https://mmsegmentation.readthedocs.io/en/latest/get_started.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Av-bYaXa_iMa",
    "outputId": "dc922937-cde2-4bed-c060-93d244dcff86"
   },
   "outputs": [],
   "source": [
    "!pip install -U openmim\n",
    "!mim install mmcv-full\n",
    "!git clone https://github.com/open-mmlab/mmsegmentation.git\n",
    "%cd mmsegmentation\n",
    "!pip install -v -e .\n",
    "# \"-v\" means verbose, or more output\n",
    "# \"-e\" means installing a project in editable mode,\n",
    "# thus any local modifications made to the code will take effect without reinstallation.\n",
    "%cd ../"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Cài đặt bổ sung một số thư viện</h3>\n",
    "</div>\n",
    "Nền tảng Google Colab cung cấp môi trường với các thư viện Machine Learning, Deep Learning cơ bản đã được cài đặt sẵn, phần này sẽ cài đặt một số thư viện sử dụng thêm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "p3VQ4ZiG6_p_",
    "outputId": "697fdc04-9e9c-4036-bf6e-eb2289428c13"
   },
   "outputs": [],
   "source": [
    "!pip install segmentation_models_pytorch\n",
    "!pip install albumentations\n",
    "!pip install --upgrade gdown"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Import thư viện</h3>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fxslSGri69QF"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import segmentation_models_pytorch as smp\n",
    "import albumentations as A\n",
    "from albumentations.pytorch import ToTensorV2 # np.array -> torch.tensor\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from glob import glob\n",
    "from torchvision.datasets import VOCSegmentation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Định nghĩa Dataset</h3>\n",
    "</div>\n",
    "Viết class kế thừa từ class Dataset cung cấp sẵn trong PyTorch để đọc dữ liệu từ ổ cứng. Yêu cầu viết đủ 3 hàm __init__() để khởi tạo class, __len__() để trả về số điểm dữ liệu có trong tập dữ liệu và __getitem__() trả về 1 điểm dữ liệu cụ thể. Trong phần này, do tập dữ liệu PASCAL VOC đã rất phổ biến nên sẽ tận dụng Class Dataset đã được viết sẵn. Tham khảo thêm: https://albumentations.ai/docs/autoalbument/examples/pascal_voc/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HQhoYMut69QG"
   },
   "outputs": [],
   "source": [
    "cv2.setNumThreads(0)\n",
    "cv2.ocl.setUseOpenCL(False)\n",
    "\n",
    "VOC_CLASSES = [\n",
    "    \"background\",\n",
    "    \"aeroplane\",\n",
    "    \"bicycle\",\n",
    "    \"bird\",\n",
    "    \"boat\",\n",
    "    \"bottle\",\n",
    "    \"bus\",\n",
    "    \"car\",\n",
    "    \"cat\",\n",
    "    \"chair\",\n",
    "    \"cow\",\n",
    "    \"diningtable\",\n",
    "    \"dog\",\n",
    "    \"horse\",\n",
    "    \"motorbike\",\n",
    "    \"person\",\n",
    "    \"potted plant\",\n",
    "    \"sheep\",\n",
    "    \"sofa\",\n",
    "    \"train\",\n",
    "    \"tv/monitor\",\n",
    "]\n",
    "\n",
    "\n",
    "VOC_COLORMAP = [\n",
    "    [0, 0, 0],\n",
    "    [128, 0, 0],\n",
    "    [0, 128, 0],\n",
    "    [128, 128, 0],\n",
    "    [0, 0, 128],\n",
    "    [128, 0, 128],\n",
    "    [0, 128, 128],\n",
    "    [128, 128, 128],\n",
    "    [64, 0, 0],\n",
    "    [192, 0, 0],\n",
    "    [64, 128, 0],\n",
    "    [192, 128, 0],\n",
    "    [64, 0, 128],\n",
    "    [192, 0, 128],\n",
    "    [64, 128, 128],\n",
    "    [192, 128, 128],\n",
    "    [0, 64, 0],\n",
    "    [128, 64, 0],\n",
    "    [0, 192, 0],\n",
    "    [128, 192, 0],\n",
    "    [0, 64, 128],\n",
    "]\n",
    "\n",
    "class PascalVOCSearchDataset(VOCSegmentation):\n",
    "    def __init__(self, root=\"~/data/pascal_voc\", image_set=\"train\", download=True, transform=None):\n",
    "        super().__init__(root=root, image_set=image_set, download=download, transform=transform)\n",
    "\n",
    "    @staticmethod\n",
    "    def _convert_to_segmentation_mask(mask):\n",
    "        # This function converts a mask from the Pascal VOC format to the format required by AutoAlbument.\n",
    "        #\n",
    "        # Pascal VOC uses an RGB image to encode the segmentation mask for that image. RGB values of a pixel\n",
    "        # encode the pixel's class.\n",
    "        #\n",
    "        # AutoAlbument requires a segmentation mask to be a NumPy array with the shape [height, width, num_classes].\n",
    "        # Each channel in this mask should encode values for a single class. Pixel in a mask channel should have\n",
    "        # a value of 1.0 if the pixel of the image belongs to this class and 0.0 otherwise.\n",
    "        height, width = mask.shape[:2]\n",
    "        segmentation_mask = np.zeros((height, width, len(VOC_COLORMAP)), dtype=np.float32)\n",
    "        for label_index, label in enumerate(VOC_COLORMAP):\n",
    "            segmentation_mask[:, :, label_index] = np.all(mask == label, axis=-1).astype(float)\n",
    "        return segmentation_mask #0, 1, 2, 3, ..., 20 (H, W, C) -> (H, W, 1) -> (H, W) #numpy\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        image = cv2.imread(self.images[index])\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        mask = cv2.imread(self.masks[index])\n",
    "        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)\n",
    "        mask = self._convert_to_segmentation_mask(mask)\n",
    "        if self.transform is not None:\n",
    "            transformed = self.transform(image=image, mask=mask)\n",
    "            image = transformed[\"image\"]\n",
    "            mask = transformed[\"mask\"]\n",
    "        return image, mask.argmax(dim=2).squeeze() #torch.tensor"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Định nghĩa các phép augmentation trên ảnh</h3>\n",
    "</div>\n",
    "Sử dụng thư viện Albumentations, tham khảo thêm: https://albumentations.ai/docs/api_reference/full_reference/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3zULyC9m69QH",
    "outputId": "975da6b3-dcda-4af1-d955-48bd05deef00"
   },
   "outputs": [],
   "source": [
    "trainsize = 256\n",
    "\n",
    "train_transform = A.Compose([\n",
    "    A.Resize(width=trainsize, height=trainsize),\n",
    "    A.HorizontalFlip(),\n",
    "    A.RandomBrightnessContrast(),\n",
    "    A.Blur(),\n",
    "    A.Sharpen(),\n",
    "    A.RGBShift(),\n",
    "    A.Cutout(num_holes=5, max_h_size=25, max_w_size=25, fill_value=0),\n",
    "    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0),\n",
    "    ToTensorV2(),\n",
    "])\n",
    "\n",
    "test_trainsform = A.Compose([\n",
    "    A.Resize(width=trainsize, height=trainsize),\n",
    "    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0),\n",
    "    ToTensorV2(), # numpy.array -> torch.tensor (B, 3, H, W)\n",
    "])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Đoạn code dùng để convert ảnh sau khi đã chuẩn hoá thành ảnh ban đầu</h3>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "akVoReo169QI"
   },
   "outputs": [],
   "source": [
    "class UnNormalize(object):\n",
    "    def __init__(self, mean, std):\n",
    "        self.mean = mean\n",
    "        self.std = std\n",
    "        \n",
    "    def __call__(self, tensor):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.\n",
    "        Returns:\n",
    "            Tensor: Normalized image.\n",
    "        \"\"\"\n",
    "        for t, m, s in zip(tensor, self.mean, self.std):\n",
    "            t.mul_(s).add_(m)\n",
    "            # The normalize code -> t.sub_(m).div_(s)\n",
    "        return tensor\n",
    "    \n",
    "unorm = UnNormalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Kiểm tra 1 cặp ảnh đầu vào và ảnh kết quả phân vùng trước khi đưa vào mô hình training</h3>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 240
    },
    "id": "5ld3Seht69QJ",
    "outputId": "9a24928d-d1d8-4206-c7e5-77c543f6da41"
   },
   "outputs": [],
   "source": [
    "train_dataset = PascalVOCSearchDataset(image_set=\"train\", download=True, transform=train_transform)\n",
    "test_dataset = PascalVOCSearchDataset(image_set=\"val\", download=False, transform=test_trainsform)\n",
    "\n",
    "image, mask = train_dataset.__getitem__(10)\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(unorm(image).permute(1, 2, 0))\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(mask)\n",
    "plt.show()\n",
    "print(mask.unique())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Tạo AverageMeter</h3>\n",
    "</div>\n",
    "AverageMeter có nhiệm vụ lưu lại giá trị trung bình của độ chính xác, giá trị hàm loss, ... trong suốt quá trình training. Tham khảo thêm: https://discuss.pytorch.org/t/meaning-of-parameters/10655"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZizwL5WZ69QJ"
   },
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Lập trình hàm tính toán độ chính xác dựa trên IoU và Dice</h3>\n",
    "</div>\n",
    "Tham khảo: https://github.com/hszhao/semseg/blob/master/util/util.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MVdLbRdF69QJ",
    "outputId": "8cba3082-fb16-4a20-c418-79b1d7e48612"
   },
   "outputs": [],
   "source": [
    "#metrics\n",
    "\n",
    "def intersectionAndUnionGPU(output, target, K, ignore_index=255):\n",
    "    # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1.\n",
    "    assert (output.dim() in [1, 2, 3])\n",
    "    assert output.shape == target.shape\n",
    "    output = output.view(-1)\n",
    "    target = target.view(-1)\n",
    "    output[target == ignore_index] = ignore_index\n",
    "    intersection = output[output == target]\n",
    "    area_intersection = torch.histc(intersection, bins=K, min=0, max=K-1)\n",
    "    area_output = torch.histc(output, bins=K, min=0, max=K-1)\n",
    "    area_target = torch.histc(target, bins=K, min=0, max=K-1)\n",
    "    area_union = area_output + area_target - area_intersection\n",
    "    return area_intersection, area_union, area_target "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Download file pretrained ImageNet của backbone MiT và convert từ sang dạng key phù hợp với lập trình của MMSegmentation</h3>\n",
    "</div>\n",
    "Tham khảo phần Usage: https://github.com/open-mmlab/mmsegmentation/tree/master/configs/segformer\n",
    "\n",
    "Tham khảo PyTorch state_dict: https://pytorch.org/tutorials/recipes/recipes/what_is_state_dict.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sPoUl5XRCCkX",
    "outputId": "07cc1ecb-0b68-42a9-8e73-2f151fc7261b"
   },
   "outputs": [],
   "source": [
    "!gdown 1EyaZVdbezIJsj8LviM7GaIBto46a1N-Z\n",
    "!gdown 1L8NYh3LOSGf7xNm7TsZVXURbYYfeJVKh\n",
    "!gdown 1m8fsG812o6KotF1NVo0YuiSfSn18TAOA\n",
    "!gdown 1d3wU8KNjPL4EqMCIEO_rO-O3-REpG82T\n",
    "!gdown 1BUtU42moYrOFbsMCE-LTTkUE-mrWnfG2\n",
    "!gdown 1d7I50jVjtCddnhpf-lqj8-f13UyCzoW1\n",
    "!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b0.pth mit_b0_mmseg.pth\n",
    "!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b1.pth mit_b1_mmseg.pth\n",
    "!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b2.pth mit_b2_mmseg.pth\n",
    "!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b3.pth mit_b3_mmseg.pth\n",
    "!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b4.pth mit_b4_mmseg.pth\n",
    "!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b5.pth mit_b5_mmseg.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Chuẩn bị cho quá trình training</h3>\n",
    "</div>\n",
    "\n",
    "\n",
    "1.   Lựa chọn device: PyTorch yêu cầu lựa chọn cụ thể device sẽ train và yêu cầu người dùng tự move dữ liệu, mô hình vào device đã lựa chọn. Device có thể là \"cuda\" - tức là GPU NVIDIA hoặc \"cpu\".\n",
    "2.   Định nghĩa DataLoader, khác với Dataset là cách đọc dữ liệu từ ổ cứng, DataLoader ghép nhiều điểm dữ liệu vào cùng nhau tạo thành 1 batch để đưa vào train mô hình. Lưu ý thêm: batch_size nên đặt là 4, 8, 16, 32, ... và nên để lớn nhất có thể\n",
    "3.   Khởi tạo mô hình. Mô hình được khởi tạo từ config (do tuân theo format của MMSegmentation). Tham khảo thêm: https://github.com/open-mmlab/mmsegmentation/blob/master/configs/_base_/models/segformer_mit-b0.py và https://github.com/open-mmlab/mmsegmentation/tree/master/configs/segformer\n",
    "4.   Khởi tạo hàm loss\n",
    "5.   Khởi tạo thuật toán tối ưu (optimizer)\n",
    "6.   Khởi tạo các độ đo sẽ sử dụng để đánh giá hiệu năng của mô hình. Phần này sẽ sử dụng các hàm độ đo Dice và IoU được lập trình sẵn trong thư viện torchmetrics\n",
    "7.   Khởi tạo từng AverageMeter để lưu lại giá trị của từng độ đo, giá trị hàm loss, thời gian train, ... trong suốt quá trình train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F9oyoO1g69QJ",
    "outputId": "651aaa43-6a67-4491-d090-0fed08af63fe"
   },
   "outputs": [],
   "source": [
    "#device\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "#load data\n",
    "batch_size = 24\n",
    "n_workers = os.cpu_count()\n",
    "print(\"num_workers =\", n_workers)\n",
    "trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,\n",
    "                                          shuffle=True, num_workers=n_workers)\n",
    "testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,\n",
    "                                          shuffle=False, num_workers=n_workers)\n",
    "\n",
    "#model settings\n",
    "norm_cfg = dict(type='BN', requires_grad=True)\n",
    "model_cfg = dict(\n",
    "    type='EncoderDecoder',\n",
    "    pretrained=None,\n",
    "    backbone=dict(\n",
    "        type='MixVisionTransformer',\n",
    "        in_channels=3,\n",
    "        embed_dims=64,\n",
    "        num_stages=4,\n",
    "        num_layers=[3, 8, 27, 3],\n",
    "        num_heads=[1, 2, 5, 8],\n",
    "        patch_sizes=[7, 3, 3, 3],\n",
    "        sr_ratios=[8, 4, 2, 1],\n",
    "        out_indices=(0, 1, 2, 3),\n",
    "        mlp_ratio=4,\n",
    "        qkv_bias=True,\n",
    "        drop_rate=0.0,\n",
    "        attn_drop_rate=0.0,\n",
    "        drop_path_rate=0.1,\n",
    "        init_cfg = dict(type=\"Pretrained\", checkpoint=\"mit_b4_mmseg.pth\")),\n",
    "    decode_head=dict(\n",
    "        type='SegformerHead',\n",
    "        in_channels=[64, 128, 320, 512],\n",
    "        in_index=[0, 1, 2, 3],\n",
    "        channels=256,\n",
    "        dropout_ratio=0.1,\n",
    "        num_classes=21,\n",
    "        norm_cfg=norm_cfg,\n",
    "        align_corners=False,\n",
    "        loss_decode=dict(\n",
    "            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),\n",
    "    # model training and testing settings\n",
    "    train_cfg=dict(),\n",
    "    test_cfg=dict(mode='whole'))\n",
    "\n",
    "import mmcv\n",
    "from mmseg.models import build_segmentor\n",
    "model = build_segmentor(model_cfg).to(device)\n",
    "model.init_weights()\n",
    "\n",
    "#loss\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "#optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)\n",
    "n_eps = 30\n",
    "\n",
    "#meter\n",
    "train_loss_meter = AverageMeter()\n",
    "intersection_meter = AverageMeter()\n",
    "union_meter = AverageMeter()\n",
    "target_meter = AverageMeter()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Training thôi ...</h3>\n",
    "</div>\n",
    "Tham khảo thêm cách viết code training trong PyTorch: https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "id": "pvVaP6fk69QK",
    "outputId": "8d6b9455-c070-49a7-e0fc-bb4cfa03a61e",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for ep in range(1, 1+n_eps):\n",
    "    train_loss_meter.reset()\n",
    "    intersection_meter.reset()\n",
    "    union_meter.reset()\n",
    "    target_meter.reset()\n",
    "    model.train()\n",
    "\n",
    "    for batch_id, (x, y) in enumerate(tqdm(trainloader), start=1):\n",
    "        optimizer.zero_grad()\n",
    "        n = x.shape[0]\n",
    "        x = x.to(device).float()\n",
    "        y = y.to(device).long()\n",
    "        y_hat = model.forward_dummy(x) #(B, C, H, W)\n",
    "        loss = criterion(y_hat, y) #(B, C, H, W) >< (B, H, W)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        #save metrics\n",
    "        with torch.no_grad():\n",
    "            train_loss_meter.update(loss.item())\n",
    "            y_hat_mask = y_hat.argmax(dim=1).squeeze(1) # (B, C, H, W) -> (B, 1, H, W) -> (B, H, W)\n",
    "            intersection, union, target = intersectionAndUnionGPU(y_hat_mask.float(), y.float(), 21)\n",
    "            intersection_meter.update(intersection)\n",
    "            union_meter.update(union)\n",
    "            target_meter.update(target)\n",
    "\n",
    "    #compute iou, dice\n",
    "    with torch.no_grad():\n",
    "        iou_class = intersection_meter.sum / (union_meter.sum + 1e-10) #vector 21D\n",
    "        dice_class = (2 * intersection_meter.sum) / (intersection_meter.sum + union_meter.sum + 1e-10) #vector 21D\n",
    "\n",
    "        mIoU = torch.mean(iou_class) #mean vector 21D\n",
    "        mDice = torch.mean(dice_class) #mean vector 21D\n",
    "\n",
    "    print(\"EP {}, current_lr = {} , train loss = {} IoU = {}, dice = {}\".format(ep, scheduler.get_last_lr(), train_loss_meter.avg, mIoU, mDice))\n",
    "    scheduler.step()\n",
    "    \n",
    "    if ep >= 25:\n",
    "        torch.save(model.state_dict(), \"modelSegFormer_ep_{}.pth\".format(ep))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Viết code đưa ra kết quả trên tập dữ liệu test</h3>\n",
    "</div>\n",
    "Lưu ý: Trong quá trình test cần chuyển model về chế độ eval và nên đặt torch.no_grad(). Điểm khác nhau giữa model.eval() và torch.no_grad() là model.eval() sẽ đặt chế độ của mô hình về chế độ evaluation (thay đổi các layer như dropout, batch norm, ... khác với quá trình training) còn torch.no_grad() sẽ tắt tính toán đạo hàm của PyTorch để tiết kiệm tính toán. Do vậy khi test mô hình nên chú ý dùng cả 2 câu lệnh này."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "9LVXDQJN69QK",
    "outputId": "b1cf73c2-8a6f-48a9-ae1b-a06612acf7be"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "test_intersection_meter = AverageMeter()\n",
    "test_union_meter = AverageMeter()\n",
    "test_target_meter = AverageMeter()\n",
    "with torch.no_grad():\n",
    "    for batch_id, (x, y) in enumerate(tqdm(testloader), start=1):\n",
    "        n = x.shape[0]\n",
    "        x = x.to(device).float()\n",
    "        y = y.to(device).long()\n",
    "        y_hat = model.forward_dummy(x)\n",
    "        y_hat = y_hat.squeeze(1)\n",
    "        y_hat_mask = y_hat.argmax(dim=1)\n",
    "        \n",
    "        intersection, union, target = intersectionAndUnionGPU(y_hat_mask, y, 21)\n",
    "        test_intersection_meter.update(intersection)\n",
    "        test_union_meter.update(union)\n",
    "        test_target_meter.update(target)\n",
    "        \n",
    "    iou_class = test_intersection_meter.sum / (test_union_meter.sum + 1e-10)\n",
    "    dice_class = 2*test_intersection_meter.sum / (test_intersection_meter.sum + test_union_meter.sum + 1e-10)\n",
    "\n",
    "    mIoU = torch.mean(iou_class)\n",
    "    mDice = torch.mean(dice_class)\n",
    "        \n",
    "print(\"TEST: IoU = {}, dice = {}\".format(mIoU, mDice))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"markdown-google-sans\">\n",
    "  <h3>Viết code hiển thị kết quả dự đoán</h3>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "lQHnQDJH69QL",
    "outputId": "5667137d-fdac-4da3-abbb-763ef19caaa7"
   },
   "outputs": [],
   "source": [
    "#predict\n",
    "import random\n",
    "id = random.randint(0, 1450)\n",
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "    x, y = test_dataset.__getitem__(id)\n",
    "    y_predict = model.forward_dummy(x.unsqueeze(0).to(device)).argmax(dim=1).squeeze().cpu().numpy()\n",
    "    # print(np.unique(y_predict))\n",
    "    # print(y_predict.shape)\n",
    "    for i in np.unique(y_predict).tolist():\n",
    "        print(VOC_CLASSES[i])\n",
    "    color_mask_predict = np.zeros((*y_predict.shape, 3))\n",
    "    for i, color in enumerate(VOC_COLORMAP):\n",
    "        color_mask_predict[y_predict==i] = np.array(color)\n",
    "    color_mask = np.zeros((*y_predict.shape, 3))\n",
    "    for i, color in enumerate(VOC_COLORMAP):\n",
    "        color_mask[y==i] = np.array(color)\n",
    "    plt.subplot(1,3,1)\n",
    "    plt.imshow(unorm(x).permute(1, 2, 0))\n",
    "    plt.subplot(1,3,2)\n",
    "    plt.imshow(color_mask)\n",
    "    plt.subplot(1,3,3)\n",
    "    plt.imshow(color_mask_predict)\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "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.8.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "ce06b65ffa354aec5ed2625f33c72f9dd46b0d2ee10cd3481aa94d5dfff5c8d3"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}