DIFL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac0d705c",
   "metadata": {},
   "source": [
    "# Importing Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bcde2de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "import sys\n",
    "import collections\n",
    "from yacs.config import CfgNode as CN\n",
    "import os\n",
    "from PIL import Image, ImageFile\n",
    "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision import transforms, datasets\n",
    "import numpy as np\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision import transforms\n",
    "from cityscapesscripts.helpers.labels import trainId2label as t2l\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from einops import rearrange\n",
    "import shutil\n",
    "import time\n",
    "import tqdm\n",
    "from csv import reader\n",
    "from functools import reduce\n",
    "import torchvision\n",
    "import torch._utils\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms.functional as TF\n",
    "import numpy as np\n",
    "import math\n",
    "from torchvision import models\n",
    "import progressbar\n",
    "from time import sleep\n",
    "from progress.spinner import MoonSpinner\n",
    "from pathlib import Path\n",
    "import csv\n",
    "import itertools\n",
    "import random\n",
    "import gc\n",
    "import statistics\n",
    "import pandas as pd\n",
    "import logging\n",
    "import functools\n",
    "class color:\n",
    "   PURPLE = '\\033[95m'\n",
    "   CYAN = '\\033[96m'\n",
    "   DARKCYAN = '\\033[36m'\n",
    "   BLUE = '\\033[94m'\n",
    "   GREEN = '\\033[92m'\n",
    "   YELLOW = '\\033[93m'\n",
    "   RED = '\\033[91m'\n",
    "   BOLD = '\\033[1m'\n",
    "   UNDERLINE = '\\033[4m'\n",
    "   END = '\\033[0m'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fffbf97c",
   "metadata": {},
   "source": [
    "# Defining Training devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c689ac8",
   "metadata": {},
   "outputs": [],
   "source": [
    "device_rtx3090 = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "device_rtx3070 = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd42e826",
   "metadata": {},
   "source": [
    "# Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e114e7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_HEIGHT, IMG_WIDTH = 224, 224\n",
    "BATCH_SIZE = 8\n",
    "NUM_EPOCHS = 500"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3654bb3c",
   "metadata": {},
   "source": [
    "# Directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "029af218",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = '/home/sourajit/DeepLearningProjects/TB/datasets'\n",
    "MODEL_PATH_GC = '/home/sourajit/DeepLearningProjects/TB/datasets/Models/'+'modelGC.pth'\n",
    "MODEL_PATH_G = '/home/sourajit/DeepLearningProjects/TB/datasets/Models/'+'modelG.pth'\n",
    "MODEL_PATH_D = '/home/sourajit/DeepLearningProjects/TB/datasets/Models/'+'modelC.pth'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df701e3d",
   "metadata": {},
   "source": [
    "# Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4d6101c",
   "metadata": {},
   "outputs": [],
   "source": [
    "domains= ['china', 'india', 'usa', 'tbx']\n",
    "domain_pairs = list(((src, tar) for src in domains for tar in domains if src != tar))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd47509",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa4a3364",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TBDataset(Dataset):\n",
    "    def __init__(self, root_dir, label_type, distribution_type, transform=None):\n",
    "        \n",
    "        self.transform = transform\n",
    "        self.label_type = label_type # 'classification_label' or 'domain_label'\n",
    "        self.distribution_type = distribution_type # 'source' or 'target'\n",
    "        \n",
    "        self.folders_list = os.listdir(root_dir)\n",
    "        self.Img_list = []\n",
    "        for folder in self.folders_list:\n",
    "            self.Img_list.extend(\n",
    "                [root_dir+'/'+folder+'/'+path for path in os.listdir(root_dir+'/'+folder)]\n",
    "            )\n",
    "        \n",
    "    def __len__(self):\n",
    "        length = len(self.Img_list)\n",
    "        return length\n",
    "      \n",
    "    def __getitem__(self, index):\n",
    "        image = Image.open(self.Img_list[index])\n",
    "        image = image.convert(\"L\")\n",
    "        label = None\n",
    "        if (self.label_type == 'classification_label'):\n",
    "            label = int(self.Img_list[index].split('/')[-2])\n",
    "        elif (self.label_type == 'domain_label'):\n",
    "            if (self.distribution_type == 'source'):\n",
    "                label = 0\n",
    "            elif (self.distribution_type == 'target'):\n",
    "                label = 1\n",
    "        if self.transform is not None:\n",
    "            image = self.transform(image)\n",
    "        x = transforms.ToTensor()(image)\n",
    "        y = np.array(label)\n",
    "        y = torch.from_numpy(y)\n",
    "        y = y.type(torch.LongTensor)\n",
    "        return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "724026cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.Resize((IMG_HEIGHT, IMG_WIDTH))\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d575268",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_subset(dataset, data_limit, num_classes=2):\n",
    "    #print (\"***1***\", data_limit)\n",
    "    data_indices = []\n",
    "    per_class_data_limit = data_limit//num_classes\n",
    "    target_counter = collections.Counter()\n",
    "    with MoonSpinner('Data Loader Working...') as bar:\n",
    "        for idx, data in enumerate(dataset):\n",
    "            if (idx == len(dataset)-1):\n",
    "                for i in range(data_limit - per_class_data_limit*num_classes):\n",
    "                    print (\"Aisi\")\n",
    "                    data_indices.append(idx)\n",
    "            target = data[1].item()\n",
    "            target_counter[target] += 1\n",
    "            if target_counter[target] <= per_class_data_limit:\n",
    "                data_indices.append(idx)\n",
    "            bar.next()\n",
    "    \n",
    "    sub_dataset = torch.utils.data.Subset(dataset, data_indices)\n",
    "    #print (\"***2***\", len(sub_dataset))\n",
    "    return sub_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b199a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stratify_split(dataset, split):\n",
    "    #split is a number between 0 to 100 train:test=split:1\n",
    "    train_indices = []\n",
    "    test_indices = []\n",
    "    target_stat = collections.Counter()\n",
    "    with MoonSpinner('Data Loader Working...') as bar:\n",
    "        for idx, data in enumerate(dataset):\n",
    "            target = data[1].item()\n",
    "            target_stat[target] += 1\n",
    "        bar.next()\n",
    "    for k in target_stat.keys():\n",
    "         target_stat[k] = int(target_stat[k]*split/100.0)\n",
    "    target_counter = collections.Counter()\n",
    "    with MoonSpinner('Data Loader Working...') as bar:\n",
    "        for idx, data in enumerate(dataset):\n",
    "            target = data[1].item()\n",
    "            target_counter[target] += 1\n",
    "            if target_counter[target] <= target_stat[target]:\n",
    "                train_indices.append(idx)\n",
    "            else:\n",
    "                test_indices.append(idx)\n",
    "        bar.next()\n",
    "    train_dataset = torch.utils.data.Subset(dataset, train_indices)\n",
    "    test_dataset = torch.utils.data.Subset(dataset, test_indices)\n",
    "    return train_dataset, test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "384d530d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_bigger_dataset(len_A, len_B):\n",
    "    if (len_A >= len_B):\n",
    "        return 2, int(len_A/len_B), len_A - len_B*int(len_A/len_B)\n",
    "    else:\n",
    "        return 1, int(len_B/len_A), len_B - len_A*int(len_B/len_A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c3dcb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_combined_TB_data(root_dir_src, root_dir_tar, \n",
    "                         transform=None, batch_size=1, split=80, \n",
    "                         shuffle=False, pin_memory=False):\n",
    "            \n",
    "    data_src_cls = TBDataset(root_dir_src, label_type='classification_label', \n",
    "                             distribution_type='None', transform=transform)\n",
    "    data_src_dom = TBDataset(root_dir_src, label_type='domain_label', \n",
    "                             distribution_type='source', transform=transform)\n",
    "    data_tar_dom = TBDataset(root_dir_tar, label_type='domain_label', \n",
    "                             distribution_type='target', transform=transform)\n",
    "    dataset_idx, multiplier, remainder = find_bigger_dataset(len(data_src_dom), \n",
    "                                                  len(data_tar_dom))\n",
    "    #print (\"Before Balancing\")\n",
    "    #print (\"dom_src, dom_tar: \", len(data_src_dom), len(data_tar_dom))\n",
    "    if (dataset_idx == 1):\n",
    "        temp = multiplier*[data_src_dom]\n",
    "        subset = create_subset(data_src_dom, remainder, num_classes=1)\n",
    "        temp.append(subset)\n",
    "        data_src_dom = torch.utils.data.ConcatDataset(temp)\n",
    "    elif (dataset_idx == 2):\n",
    "        temp = multiplier*[data_tar_dom]\n",
    "        subset = create_subset(data_tar_dom, remainder, num_classes=1)\n",
    "        temp.append(subset)\n",
    "        data_tar_dom = torch.utils.data.ConcatDataset(temp)\n",
    "    #print (\"After Balancing\")\n",
    "    #print (\"dom_src, dom_tar: \", len(data_src_dom), len(data_tar_dom))\n",
    "    combined_data_dom = torch.utils.data.ConcatDataset((data_src_dom, data_tar_dom))\n",
    "\n",
    "    dataset_idx, multiplier, remainder = find_bigger_dataset(len(data_src_cls), \n",
    "                                                  len(combined_data_dom))\n",
    "    #print (\"Before Balancing\")\n",
    "    #print (\"cls_src, dom_combined\", len(data_src_cls), len(combined_data_dom))\n",
    "    if (dataset_idx == 1):\n",
    "        temp = multiplier*[data_src_cls]\n",
    "        subset = create_subset(data_src_cls, remainder)\n",
    "        temp.append(subset)\n",
    "        data_src_cls = torch.utils.data.ConcatDataset(temp)\n",
    "    elif (dataset_idx == 2):\n",
    "        temp = multiplier*[combined_data_dom]\n",
    "        subset = create_subset(combined_data_dom, remainder)\n",
    "        temp.append(subset)\n",
    "        combined_data_dom = torch.utils.data.ConcatDataset(temp)\n",
    "    #print (\"After Balancing\")\n",
    "    #print (\"cls_src, dom_combined\", len(data_src_cls), len(combined_data_dom))\n",
    "\n",
    "    train_data_cls, test_data_cls = stratify_split(data_src_cls, split)\n",
    "    train_loader_cls = torch.utils.data.DataLoader(train_data_cls, batch_size=batch_size, \n",
    "                                              shuffle=shuffle, pin_memory=pin_memory)\n",
    "    test_loader_cls = torch.utils.data.DataLoader(test_data_cls, batch_size=batch_size, \n",
    "                                              shuffle=shuffle, pin_memory=pin_memory)\n",
    "\n",
    "    train_data_dom, test_data_dom = stratify_split(combined_data_dom, split)\n",
    "    train_loader_dom = torch.utils.data.DataLoader(train_data_dom, batch_size=batch_size, \n",
    "                                              shuffle=shuffle, pin_memory=pin_memory)\n",
    "    test_loader_dom = torch.utils.data.DataLoader(test_data_dom, batch_size=batch_size, \n",
    "                                              shuffle=shuffle, pin_memory=pin_memory)\n",
    "\n",
    "    return train_loader_cls, test_loader_cls, train_loader_dom, test_loader_dom"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79d3f9ae",
   "metadata": {},
   "source": [
    "# Utility Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93364308",
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_num_param(model):\n",
    "    return sum(p.numel() for p in model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e94a110",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_param(num_param):\n",
    "    num_param = str(num_param)\n",
    "    ret = []\n",
    "    length = len(num_param)\n",
    "    for i in range(length):\n",
    "        idx = i+1\n",
    "        if (idx>1 and (idx-1)%3==0):\n",
    "            ret.append(',')\n",
    "        ret.append(num_param[length-idx])\n",
    "    ret.reverse()\n",
    "    ret = ''.join(ret)\n",
    "    temp_a = '## Number of parameters: ' + ret  + ' ##'\n",
    "    temp_b = '#' * len(temp_a)\n",
    "    columns = shutil.get_terminal_size().columns\n",
    "    print (f'{temp_b}'.center(columns))\n",
    "    print (f'{temp_a}'.center(columns))\n",
    "    print (f'{temp_b}'.center(columns))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e7562f6",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3e72418",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(torch.nn.Module):\n",
    "    def __init__(self, num_of_filters):\n",
    "        super(Generator, self).__init__()\n",
    "        self.num_of_filters = num_of_filters\n",
    "        self.Resnet50_model = models.resnet50(pretrained=False)\n",
    "        self.Resnet50_model = nn.Sequential(*list(self.Resnet50_model.children())[:-2])\n",
    "        self.Resnet50_model[0] = nn.Conv2d(1,64,7,2,3)\n",
    "        self.Resnet50_model = nn.Sequential(*list(self.Resnet50_model.children()))\n",
    "        \n",
    "        self.A_UpConv_1 = nn.ConvTranspose2d(2048, 1024, (3,3), stride=(2,2))\n",
    "        self.A_activation_1 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.A_UpConv_2 = nn.ConvTranspose2d(1024, self.num_of_filters, (3,3), \n",
    "                                             stride=(2,2))\n",
    "        self.A_activation_2 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.A_conv = nn.Conv2d(self.num_of_filters, self.num_of_filters, \n",
    "                                kernel_size=3, bias=False)\n",
    "        self.A_activation_3 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.A_maxpool = nn.MaxPool2d(3, stride=2)\n",
    "        \n",
    "        self.B_UpConv_1 = nn.ConvTranspose2d(self.num_of_filters, self.num_of_filters, \n",
    "                                             (3,3), stride=(2,2))\n",
    "        self.B_activation_1 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.B_conv = nn.Conv2d(self.num_of_filters, self.num_of_filters, \n",
    "                                kernel_size=3, bias=False)\n",
    "        self.B_activation_3 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        #self.B_maxpool = nn.MaxPool2d(3, stride=2)\n",
    "        \n",
    "        self.C_UpConv_1 = nn.ConvTranspose2d(self.num_of_filters, self.num_of_filters, \n",
    "                                             (3,3), stride=(2,2))\n",
    "        self.C_activation_1 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.Resnet50_model(x)\n",
    "        x = self.A_UpConv_1(x)\n",
    "        x = self.A_activation_1(x)\n",
    "        x = self.A_UpConv_2(x)\n",
    "        x = self.A_activation_2(x)\n",
    "        x = self.A_conv(x)\n",
    "        x = self.A_activation_3(x)\n",
    "        x = self.A_maxpool(x)\n",
    "        x = self.B_UpConv_1(x)\n",
    "        x = self.B_activation_1(x)\n",
    "        x = self.B_conv(x)\n",
    "        x = self.B_activation_3(x)\n",
    "        #x = self.B_maxpool(x)\n",
    "        x = self.C_UpConv_1(x)\n",
    "        x = self.C_activation_1(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb59f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(torch.nn.Module):\n",
    "    def __init__(self, num_of_filters):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.num_of_filters = num_of_filters\n",
    "        self.vgg19_model = models.vgg19(pretrained=False)\n",
    "        self.vgg19_model.features[0] = nn.Conv2d(512,64,3,1,1)\n",
    "        self.vgg19_model = nn.Sequential(*list(self.vgg19_model.children())[:-2])\n",
    "        self.Linear_1 = nn.Linear(self.num_of_filters, 256)\n",
    "        self.Activation_1 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.Linear_2 = nn.Linear(256, 128)\n",
    "        self.Activation_2 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.Linear_3 = nn.Linear(128, 64)\n",
    "        self.Activation_3 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.Linear_final = nn.Linear(64, 2)\n",
    "        self.Activation_final = nn.Sigmoid()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.vgg19_model(x)\n",
    "        x = nn.Flatten()(x)\n",
    "        x = self.Linear_1(x)\n",
    "        x = self.Activation_1(x)\n",
    "        x = self.Linear_2(x)\n",
    "        x = self.Activation_2(x)\n",
    "        x = self.Linear_3(x)\n",
    "        x = self.Activation_3(x)\n",
    "        x = self.Linear_final(x)\n",
    "        x = self.Activation_final(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a175d353",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Classifier(torch.nn.Module):\n",
    "    def __init__(self, num_of_filters):\n",
    "        super(Classifier, self).__init__()\n",
    "        self.num_of_filters = num_of_filters\n",
    "        self.vgg19_model = models.vgg19(pretrained=False)\n",
    "        self.vgg19_model.features[0] = nn.Conv2d(512,64,3,1,1)\n",
    "        self.vgg19_model = nn.Sequential(*list(self.vgg19_model.children())[:-2])\n",
    "        self.Linear_1 = nn.Linear(self.num_of_filters, 256)\n",
    "        self.Activation_1 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.Linear_2 = nn.Linear(256, 128)\n",
    "        self.Activation_2 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.Linear_3 = nn.Linear(128, 64)\n",
    "        self.Activation_3 = torch.nn.LeakyReLU(negative_slope=0.2, inplace=False)\n",
    "        self.Linear_final = nn.Linear(64, 2)\n",
    "        self.Activation_final = nn.Sigmoid()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.vgg19_model(x)\n",
    "        x = nn.Flatten()(x)\n",
    "        x = self.Linear_1(x)\n",
    "        x = self.Activation_1(x)\n",
    "        x = self.Linear_2(x)\n",
    "        x = self.Activation_2(x)\n",
    "        x = self.Linear_3(x)\n",
    "        x = self.Activation_3(x)\n",
    "        x = self.Linear_final(x)\n",
    "        x = self.Activation_final(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecc1c95c",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = Generator(512)\n",
    "#print_param(count_num_param(G))\n",
    "D = Discriminator(512)\n",
    "#print_param(count_num_param(D))\n",
    "C = Classifier(512)\n",
    "#print_param(count_num_param(C))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c363d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.BCELoss()\n",
    "Params_Classification = list(G.parameters()) + list(C.parameters())\n",
    "Optimizer_Classification = optim.Adam(Params_Classification, \n",
    "                                      lr=0.0002, betas=(0.5, 0.999))\n",
    "Params_DomainInvariance = list(G.parameters()) + list(D.parameters())\n",
    "Optimizer_DomainInvariance_G = optim.Adam(G.parameters(), \n",
    "                                        lr=0.0002, betas=(0.5, 0.999))\n",
    "Optimizer_DomainInvariance_D = optim.Adam(D.parameters(), \n",
    "                                        lr=0.0002, betas=(0.5, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d4897c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def training(diseases_classification_trainset, domain_classification_trainset, \n",
    "            MODEL_PATH_GC, MODEL_PATH_G, MODEL_PATH_D, device, num_epochs):\n",
    "    # 1st step: Train (DIF Generator -> disseases classifier)\n",
    "    # 2nd step: Train (DIF Generator -> domain classifier)\n",
    "    G.to(device)\n",
    "    C.to(device)\n",
    "    D.to(device)\n",
    "    G.eval()\n",
    "    C.eval()\n",
    "    D.eval()\n",
    "    for epoch in range(num_epochs):\n",
    "        # Training The Classifier with (X_s, y_s) : Class Labels\n",
    "        for batch_number, example in enumerate(diseases_classification_trainset):\n",
    "            X_s, y_s = example\n",
    "            X_s, y_s = X_s.to(device), y_s.to(device)\n",
    "            G_X_s = G(X_s)\n",
    "            y_hat_s = C(G_X_s)\n",
    "            l_c = criterion(y_hat_s, y_s)\n",
    "            Optimizer_Classification.zero_grad()\n",
    "            l_c.backward()\n",
    "            Optimizer_Classification.step()\n",
    "        torch.save({\n",
    "                'model_state_dict': model.state_dict(),\n",
    "                'optim_state_dict': Optimizer_Classification.state_dict(),\n",
    "                'epoch': epoch\n",
    "        }, MODEL_PATH_GC)\n",
    "        # Training The Domain Invariance Step with (X_s+t=X, d_s+t=d) : Domain Labels\n",
    "        for batch_number, example in enumerate(domain_classification_trainset):\n",
    "            X, d = example\n",
    "            X, d = X.to(device), d.to(device)\n",
    "            G_X = G(X)\n",
    "            d_hat = D(G_X)\n",
    "            l_d = criterion(d_hat, d)\n",
    "            d_gen = torch.full(d.shape, 0.5)\n",
    "            d_gen.to(device)\n",
    "            l_g = criterion(d_hat, d_gen)\n",
    "            \n",
    "            Optimizer_DomainInvariance_G.zero_grad()\n",
    "            Optimizer_DomainInvariance_D.zero_grad()\n",
    "            l_d.backward()\n",
    "            l_g.backward()\n",
    "            Optimizer_DomainInvariance_G.step()\n",
    "            Optimizer_DomainInvariance_D.step()\n",
    "        torch.save({\n",
    "                'model_state_dict': model.state_dict(),\n",
    "                'optim_state_dict': Optimizer_DomainInvariance_G.state_dict(),\n",
    "                'epoch': epoch\n",
    "        }, MODEL_PATH_G)\n",
    "        torch.save({\n",
    "                'model_state_dict': model.state_dict(),\n",
    "                'optim_state_dict': Optimizer_DomainInvariance_D.state_dict(),\n",
    "                'epoch': epoch\n",
    "        }, MODEL_PATH_D)\n",
    "    return"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3d4f2ad",
   "metadata": {},
   "source": [
    "# Training The model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cb29cc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(domain_pairs)):\n",
    "    src_dir = data_dir+'/'+domain_pairs[i][0]\n",
    "    tar_dir = data_dir+'/'+domain_pairs[i][1]\n",
    "    diseases_classification_trainset, diseases_classification_testset, \\\n",
    "    domain_classification_trainset, domain_classification_testset = \\\n",
    "    get_combined_TB_data(root_dir_src=src_dir, root_dir_tar=tar_dir, \n",
    "                         transform=transform, split=80, batch_size=BATCH_SIZE, \n",
    "                         shuffle=True, pin_memory=False)\n",
    "    training(diseases_classification_trainset, domain_classification_trainset, \n",
    "            MODEL_PATH_GC, MODEL_PATH_G, MODEL_PATH_D, device, NUM_EPOCHS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aee5d1c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "221932bb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79985530",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a6c921b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a416f11a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}