Adversarial_Network_Compression.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/leonsuarez24/Notebooks/blob/main/Adversarial_Network_Compression.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": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CL65TKqxWLfJ",
        "outputId": "e181c4d4-86f3-4c22-e00a-6712be3f03dc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive/\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-P3vIEfqWDbF",
        "outputId": "016b0e5f-1208-4fec-c196-41c9303f9964"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-2-15c2c979f39a>:13: DeprecationWarning: Please use `rotate` from the `scipy.ndimage` namespace, the `scipy.ndimage.interpolation` namespace is deprecated.\n",
            "  from scipy.ndimage.interpolation import rotate as scipyrotate\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting torchinfo\n",
            "  Downloading torchinfo-1.8.0-py3-none-any.whl (23 kB)\n",
            "Installing collected packages: torchinfo\n",
            "Successfully installed torchinfo-1.8.0\n",
            "Collecting torchmetrics\n",
            "  Downloading torchmetrics-1.0.2-py3-none-any.whl (731 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m731.1/731.1 kB\u001b[0m \u001b[31m8.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy>1.20.0 in /usr/local/lib/python3.10/dist-packages (from torchmetrics) (1.22.4)\n",
            "Requirement already satisfied: torch>=1.8.1 in /usr/local/lib/python3.10/dist-packages (from torchmetrics) (2.0.1+cu118)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from torchmetrics) (23.1)\n",
            "Collecting lightning-utilities>=0.7.0 (from torchmetrics)\n",
            "  Downloading lightning_utilities-0.9.0-py3-none-any.whl (23 kB)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from lightning-utilities>=0.7.0->torchmetrics) (4.7.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (3.12.2)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (1.11.1)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (3.1.2)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.8.1->torchmetrics) (2.0.0)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.8.1->torchmetrics) (3.25.2)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.8.1->torchmetrics) (16.0.6)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.8.1->torchmetrics) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.8.1->torchmetrics) (1.3.0)\n",
            "Installing collected packages: lightning-utilities, torchmetrics\n",
            "Successfully installed lightning-utilities-0.9.0 torchmetrics-1.0.2\n"
          ]
        }
      ],
      "source": [
        "import sys\n",
        "import os\n",
        "import time\n",
        "import copy\n",
        "import argparse\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torchvision.utils import save_image\n",
        "from torch.utils.data import Dataset\n",
        "from torchvision import datasets, transforms\n",
        "from scipy.ndimage.interpolation import rotate as scipyrotate\n",
        "!pip install torchinfo\n",
        "from torchinfo import summary\n",
        "from torch.autograd import Function\n",
        "import torchvision.utils as vutils\n",
        "from torch.optim.lr_scheduler import StepLR\n",
        "from tqdm import tqdm\n",
        "import torchvision.utils as vutils\n",
        "!pip install torchmetrics\n",
        "from torchmetrics import Accuracy\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "import matplotlib.pyplot as plt\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8mEOqPQlWRsn"
      },
      "source": [
        "# Utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "LtD9pee3WRDr"
      },
      "outputs": [],
      "source": [
        "class AverageMeter(object):\n",
        "    \"\"\"Computes and stores the average and current value\"\"\"\n",
        "\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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "HN2XhgnEbuno"
      },
      "outputs": [],
      "source": [
        "def get_time():\n",
        "    return str(time.strftime(\"[%Y-%m-%d %H:%M:%S]\", time.localtime()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "01btqAhhWTKK"
      },
      "source": [
        "# Networks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_DOuOpwoWW7Y"
      },
      "source": [
        "##Teacher"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "dgJ65Pz5WWVY"
      },
      "outputs": [],
      "source": [
        "class VGG(nn.Module):\n",
        "    def __init__(self, vgg_name):\n",
        "        super(VGG, self).__init__()\n",
        "        self.features = self._make_layers(cfg[vgg_name])\n",
        "        self.fc = nn.Linear(512,84)\n",
        "        self.classifier = nn.Linear(84, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.features(x)\n",
        "        out = out.view(out.size(0), -1)\n",
        "        features = self.fc(out)\n",
        "        logits = self.classifier(features)\n",
        "        return logits, features\n",
        "\n",
        "    def _make_layers(self, cfg):\n",
        "        layers = []\n",
        "        in_channels = 3\n",
        "        for x in cfg:\n",
        "            if x == 'M':\n",
        "                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]\n",
        "            else:\n",
        "                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),\n",
        "                           nn.BatchNorm2d(x),\n",
        "                           nn.ReLU(inplace=True)]\n",
        "                in_channels = x\n",
        "        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "cfg = {\n",
        "    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
        "    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
        "    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],\n",
        "    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lG-kBmFQ1HCx",
        "outputId": "1c0852b3-25ce-4cf6-becf-1b8e43cd24e9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "==========================================================================================\n",
            "Layer (type:depth-idx)                   Output Shape              Param #\n",
            "==========================================================================================\n",
            "VGG                                      [12, 10]                  --\n",
            "├─Sequential: 1-1                        [12, 512, 1, 1]           --\n",
            "│    └─Conv2d: 2-1                       [12, 64, 32, 32]          1,792\n",
            "│    └─BatchNorm2d: 2-2                  [12, 64, 32, 32]          128\n",
            "│    └─ReLU: 2-3                         [12, 64, 32, 32]          --\n",
            "│    └─MaxPool2d: 2-4                    [12, 64, 16, 16]          --\n",
            "│    └─Conv2d: 2-5                       [12, 128, 16, 16]         73,856\n",
            "│    └─BatchNorm2d: 2-6                  [12, 128, 16, 16]         256\n",
            "│    └─ReLU: 2-7                         [12, 128, 16, 16]         --\n",
            "│    └─MaxPool2d: 2-8                    [12, 128, 8, 8]           --\n",
            "│    └─Conv2d: 2-9                       [12, 256, 8, 8]           295,168\n",
            "│    └─BatchNorm2d: 2-10                 [12, 256, 8, 8]           512\n",
            "│    └─ReLU: 2-11                        [12, 256, 8, 8]           --\n",
            "│    └─Conv2d: 2-12                      [12, 256, 8, 8]           590,080\n",
            "│    └─BatchNorm2d: 2-13                 [12, 256, 8, 8]           512\n",
            "│    └─ReLU: 2-14                        [12, 256, 8, 8]           --\n",
            "│    └─MaxPool2d: 2-15                   [12, 256, 4, 4]           --\n",
            "│    └─Conv2d: 2-16                      [12, 512, 4, 4]           1,180,160\n",
            "│    └─BatchNorm2d: 2-17                 [12, 512, 4, 4]           1,024\n",
            "│    └─ReLU: 2-18                        [12, 512, 4, 4]           --\n",
            "│    └─Conv2d: 2-19                      [12, 512, 4, 4]           2,359,808\n",
            "│    └─BatchNorm2d: 2-20                 [12, 512, 4, 4]           1,024\n",
            "│    └─ReLU: 2-21                        [12, 512, 4, 4]           --\n",
            "│    └─MaxPool2d: 2-22                   [12, 512, 2, 2]           --\n",
            "│    └─Conv2d: 2-23                      [12, 512, 2, 2]           2,359,808\n",
            "│    └─BatchNorm2d: 2-24                 [12, 512, 2, 2]           1,024\n",
            "│    └─ReLU: 2-25                        [12, 512, 2, 2]           --\n",
            "│    └─Conv2d: 2-26                      [12, 512, 2, 2]           2,359,808\n",
            "│    └─BatchNorm2d: 2-27                 [12, 512, 2, 2]           1,024\n",
            "│    └─ReLU: 2-28                        [12, 512, 2, 2]           --\n",
            "│    └─MaxPool2d: 2-29                   [12, 512, 1, 1]           --\n",
            "│    └─AvgPool2d: 2-30                   [12, 512, 1, 1]           --\n",
            "├─Linear: 1-2                            [12, 84]                  43,092\n",
            "├─Linear: 1-3                            [12, 10]                  850\n",
            "==========================================================================================\n",
            "Total params: 9,269,926\n",
            "Trainable params: 9,269,926\n",
            "Non-trainable params: 0\n",
            "Total mult-adds (G): 1.84\n",
            "==========================================================================================\n",
            "Input size (MB): 0.15\n",
            "Forward/backward pass size (MB): 29.11\n",
            "Params size (MB): 37.08\n",
            "Estimated Total Size (MB): 66.33\n",
            "==========================================================================================\n"
          ]
        }
      ],
      "source": [
        "t = VGG('VGG11')\n",
        "print(summary(t, input_size=(12, 3, 32, 32)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0KgbVA5xWYWR"
      },
      "source": [
        "## Student"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "JgOvVjFuWZim"
      },
      "outputs": [],
      "source": [
        "class StudentNetwork(nn.Module):\n",
        "    def __init__(self, channel, num_classes):\n",
        "        super(StudentNetwork, self).__init__()\n",
        "        self.features = nn.Sequential(\n",
        "            nn.Conv2d(channel, 6, kernel_size=5, padding=2 if channel==1 else 0),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
        "            nn.Conv2d(6, 16, kernel_size=5),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
        "        )\n",
        "        self.fc_1 = nn.Linear(16 * 5 * 5, 120)\n",
        "        self.fc_2 = nn.Linear(120, 84)\n",
        "        self.fc_3 = nn.Linear(84, num_classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.features(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = F.relu(self.fc_1(x))\n",
        "        features = F.relu(self.fc_2(x))\n",
        "        logits = self.fc_3(features)\n",
        "        return logits, features"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KeMHktxO29b7",
        "outputId": "266320f8-9504-458f-c66b-954c37b8db89"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "==========================================================================================\n",
            "Layer (type:depth-idx)                   Output Shape              Param #\n",
            "==========================================================================================\n",
            "StudentNetwork                           [12, 10]                  --\n",
            "├─Sequential: 1-1                        [12, 16, 5, 5]            --\n",
            "│    └─Conv2d: 2-1                       [12, 6, 28, 28]           456\n",
            "│    └─ReLU: 2-2                         [12, 6, 28, 28]           --\n",
            "│    └─MaxPool2d: 2-3                    [12, 6, 14, 14]           --\n",
            "│    └─Conv2d: 2-4                       [12, 16, 10, 10]          2,416\n",
            "│    └─ReLU: 2-5                         [12, 16, 10, 10]          --\n",
            "│    └─MaxPool2d: 2-6                    [12, 16, 5, 5]            --\n",
            "├─Linear: 1-2                            [12, 120]                 48,120\n",
            "├─Linear: 1-3                            [12, 84]                  10,164\n",
            "├─Linear: 1-4                            [12, 10]                  850\n",
            "==========================================================================================\n",
            "Total params: 62,006\n",
            "Trainable params: 62,006\n",
            "Non-trainable params: 0\n",
            "Total mult-adds (M): 7.90\n",
            "==========================================================================================\n",
            "Input size (MB): 0.15\n",
            "Forward/backward pass size (MB): 0.63\n",
            "Params size (MB): 0.25\n",
            "Estimated Total Size (MB): 1.02\n",
            "==========================================================================================\n"
          ]
        }
      ],
      "source": [
        "s = StudentNetwork(3,10)\n",
        "print(summary(s, input_size=(12, 3, 32, 32)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RxNlwvfeaE3H"
      },
      "source": [
        "## Discriminator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "nDeAxQIraGOd"
      },
      "outputs": [],
      "source": [
        "class Discriminator(nn.Module):\n",
        "    def __init__(self, input):\n",
        "        super(Discriminator, self).__init__()\n",
        "        self.fc1 = nn.Linear(input, 128)\n",
        "        self.fc2 = nn.Linear(128, 256)\n",
        "        self.fc3 = nn.Linear(256, 128)\n",
        "\n",
        "        self.dropout = nn.Dropout(p=0.5)\n",
        "        self.output_layer = nn.Linear(128, 1)\n",
        "        self.sigmoid = nn.Sigmoid()\n",
        "\n",
        "    def forward(self, x, drop):\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = F.relu(self.fc3(x))\n",
        "        if drop == True:\n",
        "            x = self.dropout(x)\n",
        "        x = self.sigmoid(self.output_layer(x))\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "p7YuH7ULwxPf",
        "outputId": "646f5015-e40f-4e7c-d8ca-998e8ca569a3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([12, 1])\n"
          ]
        }
      ],
      "source": [
        "model = Discriminator(100)\n",
        "input_data = torch.randn(12,100)  # Example input with shape [batch_size, channels, height, width]\n",
        "output = model(input_data, True)\n",
        "print(output.shape)  # Should be [12, 10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eZnxcQdrah-e"
      },
      "source": [
        "# Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "PoLDhxSzaijk"
      },
      "outputs": [],
      "source": [
        "def get_dataset(dataset, data_path):\n",
        "\n",
        "    if dataset == 'MNIST':\n",
        "        channel = 1\n",
        "        im_size = (28, 28)\n",
        "        num_classes = 10\n",
        "        transform = transforms.Compose([transforms.ToTensor()])\n",
        "        dst_train = datasets.MNIST(data_path, train=True, download=True, transform=transform) # no augmentation\n",
        "        dst_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)\n",
        "        class_names = [str(c) for c in range(num_classes)]\n",
        "\n",
        "    elif dataset == 'FashionMNIST':\n",
        "        channel = 1\n",
        "        im_size = (28, 28)\n",
        "        num_classes = 10\n",
        "        transform = transforms.Compose([transforms.ToTensor()])\n",
        "        dst_train = datasets.FashionMNIST(data_path, train=True, download=True, transform=transform) # no augmentation\n",
        "        dst_test = datasets.FashionMNIST(data_path, train=False, download=True, transform=transform)\n",
        "        class_names = dst_train.classes\n",
        "\n",
        "    elif dataset == 'CIFAR10':\n",
        "        channel = 3\n",
        "        im_size = (32, 32)\n",
        "        num_classes = 10\n",
        "        transform = transforms.Compose([transforms.ToTensor()])\n",
        "        #transform = transforms.Compose([transforms.ToTensor(), transforms.RandomHorizontalFlip(0.5)])\n",
        "        dst_train = datasets.CIFAR10(data_path, train=True, download=True, transform=transform) # no augmentation\n",
        "        dst_test = datasets.CIFAR10(data_path, train=False, download=True, transform=transform)\n",
        "        class_names = dst_train.classes\n",
        "\n",
        "    else:\n",
        "        exit('unknown dataset: %s'%dataset)\n",
        "\n",
        "    testloader = torch.utils.data.DataLoader(dst_test, batch_size=32, shuffle=False, num_workers=0)\n",
        "    trainloader = torch.utils.data.DataLoader(dst_train, batch_size=32, shuffle=True, num_workers=0)\n",
        "    return channel, im_size, num_classes, class_names, dst_train, dst_test, testloader, trainloader"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1p-I8b0Rbmus"
      },
      "source": [
        "# Configs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "vuzJhu4Fbnv_"
      },
      "outputs": [],
      "source": [
        "save_path = f'/content/drive/MyDrive/Proyectos/Adversarial Knowledge Distillation/experiments/{get_time()}'\n",
        "\n",
        "tb_path_teacher = save_path + '/tensorboard_teacher'\n",
        "tb_path_student = save_path + '/tensorboard_student'\n",
        "tb_path_student_distilled = save_path + '/tensorboard_student_distilled'\n",
        "\n",
        "model_teacher = save_path + '/model_teacher'\n",
        "model_student = save_path + '/model_student'\n",
        "model_student_distilled = save_path + '/model_student_distilled'\n",
        "\n",
        "os.makedirs(save_path, exist_ok=True)\n",
        "os.makedirs(tb_path_teacher, exist_ok=True)\n",
        "os.makedirs(tb_path_student, exist_ok=True)\n",
        "os.makedirs(model_teacher, exist_ok=True)\n",
        "os.makedirs(model_student, exist_ok=True)\n",
        "os.makedirs(tb_path_student_distilled, exist_ok=True)\n",
        "os.makedirs(model_student_distilled, exist_ok=True)\n",
        "\n",
        "\n",
        "\n",
        "writer_teacher = SummaryWriter(tb_path_teacher)\n",
        "writer_student = SummaryWriter(tb_path_student)\n",
        "writer_student_distilled = SummaryWriter(tb_path_student_distilled)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gjWxDO8tcV7q",
        "outputId": "a2b31789-9f1c-402e-c95b-36fae295d674"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to data/cifar-10-python.tar.gz\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 170498071/170498071 [00:13<00:00, 12855416.13it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting data/cifar-10-python.tar.gz to data\n",
            "Files already downloaded and verified\n"
          ]
        }
      ],
      "source": [
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "lr_net = 0.001\n",
        "teacher_epochs = 20\n",
        "student_epochs = 10\n",
        "dataset = 'CIFAR10'\n",
        "data_path = 'data'\n",
        "channel, im_size, num_classes, class_names, dst_train, dst_test, testloader, trainloader = get_dataset(dataset, data_path)\n",
        "accuracy = Accuracy(task=\"multiclass\", num_classes=num_classes, top_k=1).to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "utHZq_6QbYtv"
      },
      "source": [
        "# Baseline"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "teacher = VGG('VGG11').to(device)\n",
        "optimizer = torch.optim.Adam(teacher.parameters(), lr=lr_net)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "\n",
        "for epoch in range(teacher_epochs):\n",
        "\n",
        "    train_loss = AverageMeter()\n",
        "    accuracy_train = AverageMeter()\n",
        "    data_loop_train = tqdm(enumerate(trainloader), total=len(trainloader), colour='red')\n",
        "\n",
        "    teacher.train(True)\n",
        "    for _, train_data in data_loop_train:\n",
        "        train_img, train_label = train_data\n",
        "        train_img = train_img.to(device)\n",
        "        train_label = train_label.to(device)\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        train_pred, _ = teacher(train_img)\n",
        "        loss = criterion(train_pred, train_label)\n",
        "        acc = accuracy(train_pred, train_label)\n",
        "\n",
        "        train_loss.update(loss.item(), train_img.size(0))\n",
        "        accuracy_train.update(acc.item(), train_img.size(0))\n",
        "\n",
        "        dict_metrics = dict(loss = train_loss.avg, acc = accuracy_train.avg)\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        data_loop_train.set_description(f'Train Epoch [{epoch + 1} / {teacher_epochs}]')\n",
        "        data_loop_train.set_postfix(**dict_metrics)\n",
        "\n",
        "        for key, value in dict_metrics.items():\n",
        "              writer_teacher.add_scalar(f'train_{key}', value, epoch)\n",
        "\n",
        "\n",
        "    # Evaluation phase\n",
        "    teacher.eval()\n",
        "    data_loop_test = tqdm(enumerate(testloader), total=len(testloader),colour='green')\n",
        "    with torch.no_grad():\n",
        "\n",
        "        test_loss = AverageMeter()\n",
        "        accuracy_test = AverageMeter()\n",
        "\n",
        "        for _, test_data in data_loop_test:\n",
        "            test_img, test_label = test_data\n",
        "            test_img = test_img.to(device)\n",
        "            test_label = test_label.to(device)\n",
        "\n",
        "            test_pred, _ = teacher(test_img)\n",
        "            loss = criterion(test_pred, test_label)\n",
        "            acc = accuracy(test_pred, test_label)\n",
        "\n",
        "            test_loss.update(loss.item(), test_img.size(0))\n",
        "            accuracy_test.update(acc.item(), test_img.size(0))\n",
        "            dict_metrics = dict(loss = test_loss.avg, acc = accuracy_test.avg)\n",
        "\n",
        "            data_loop_test.set_description(f'Test  Epoch [{epoch + 1} / {teacher_epochs}]')\n",
        "            data_loop_test.set_postfix(**dict_metrics)\n",
        "\n",
        "            for key, value in dict_metrics.items():\n",
        "                writer_teacher.add_scalar(f'test_{key}', value, epoch)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HcTrQBvIUwFW",
        "outputId": "b1203a80-d766-42a9-901c-e8f4be3f68ba"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Train Epoch [1 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:37<00:00, 41.41it/s, acc=0.508, loss=1.34]\n",
            "Test  Epoch [1 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 56.62it/s, acc=0.671, loss=0.927]\n",
            "Train Epoch [2 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 43.06it/s, acc=0.687, loss=0.896]\n",
            "Test  Epoch [2 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 51.15it/s, acc=0.652, loss=1.02]\n",
            "Train Epoch [3 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:35<00:00, 44.45it/s, acc=0.763, loss=0.689]\n",
            "Test  Epoch [3 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 50.23it/s, acc=0.743, loss=0.753]\n",
            "Train Epoch [4 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:34<00:00, 44.98it/s, acc=0.808, loss=0.561]\n",
            "Test  Epoch [4 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 56.76it/s, acc=0.773, loss=0.69]\n",
            "Train Epoch [5 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:37<00:00, 42.14it/s, acc=0.844, loss=0.454]\n",
            "Test  Epoch [5 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 50.83it/s, acc=0.78, loss=0.681]\n",
            "Train Epoch [6 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:35<00:00, 44.17it/s, acc=0.88, loss=0.355]\n",
            "Test  Epoch [6 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 47.91it/s, acc=0.797, loss=0.653]\n",
            "Train Epoch [7 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:34<00:00, 44.99it/s, acc=0.906, loss=0.276]\n",
            "Test  Epoch [7 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 56.73it/s, acc=0.821, loss=0.599]\n",
            "Train Epoch [8 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 43.31it/s, acc=0.929, loss=0.211]\n",
            "Test  Epoch [8 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 58.39it/s, acc=0.825, loss=0.591]\n",
            "Train Epoch [9 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:35<00:00, 44.14it/s, acc=0.944, loss=0.165]\n",
            "Test  Epoch [9 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 50.95it/s, acc=0.817, loss=0.627]\n",
            "Train Epoch [10 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 43.20it/s, acc=0.955, loss=0.129]\n",
            "Test  Epoch [10 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 48.95it/s, acc=0.82, loss=0.671]\n",
            "Train Epoch [11 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:34<00:00, 44.97it/s, acc=0.965, loss=0.105]\n",
            "Test  Epoch [11 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 57.57it/s, acc=0.826, loss=0.69]\n",
            "Train Epoch [12 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 43.26it/s, acc=0.969, loss=0.0897]\n",
            "Test  Epoch [12 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 56.70it/s, acc=0.829, loss=0.703]\n",
            "Train Epoch [13 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:34<00:00, 44.83it/s, acc=0.974, loss=0.0794]\n",
            "Test  Epoch [13 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 49.43it/s, acc=0.82, loss=0.753]\n",
            "Train Epoch [14 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 43.18it/s, acc=0.978, loss=0.0649]\n",
            "Test  Epoch [14 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 51.55it/s, acc=0.812, loss=0.83]\n",
            "Train Epoch [15 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:35<00:00, 44.13it/s, acc=0.977, loss=0.068]\n",
            "Test  Epoch [15 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 57.27it/s, acc=0.815, loss=0.766]\n",
            "Train Epoch [16 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 43.06it/s, acc=0.982, loss=0.055]\n",
            "Test  Epoch [16 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 56.83it/s, acc=0.798, loss=1]\n",
            "Train Epoch [17 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:35<00:00, 44.65it/s, acc=0.982, loss=0.0551]\n",
            "Test  Epoch [17 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 51.63it/s, acc=0.832, loss=0.767]\n",
            "Train Epoch [18 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 42.80it/s, acc=0.984, loss=0.0486]\n",
            "Test  Epoch [18 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 57.87it/s, acc=0.82, loss=0.852]\n",
            "Train Epoch [19 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:36<00:00, 42.43it/s, acc=0.984, loss=0.0494]\n",
            "Test  Epoch [19 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 56.82it/s, acc=0.827, loss=0.825]\n",
            "Train Epoch [20 / 20]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:35<00:00, 43.43it/s, acc=0.986, loss=0.0434]\n",
            "Test  Epoch [20 / 20]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:06<00:00, 48.92it/s, acc=0.818, loss=0.84]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z-M23-6cl89v",
        "outputId": "f08f2da7-f443-420f-cf95-ddc8f601d098"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Train Epoch [1 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:27<00:00, 57.69it/s, acc=0.354, loss=1.75]\n",
            "Test  Epoch [1 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 71.19it/s, acc=0.424, loss=1.57]\n",
            "Train Epoch [2 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:27<00:00, 57.86it/s, acc=0.466, loss=1.46]\n",
            "Test  Epoch [2 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 71.26it/s, acc=0.49, loss=1.4]\n",
            "Train Epoch [3 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:27<00:00, 56.38it/s, acc=0.509, loss=1.36]\n",
            "Test  Epoch [3 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 61.45it/s, acc=0.504, loss=1.38]\n",
            "Train Epoch [4 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:26<00:00, 58.71it/s, acc=0.534, loss=1.29]\n",
            "Test  Epoch [4 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 63.61it/s, acc=0.521, loss=1.32]\n",
            "Train Epoch [5 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:26<00:00, 59.45it/s, acc=0.551, loss=1.25]\n",
            "Test  Epoch [5 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 61.05it/s, acc=0.546, loss=1.27]\n",
            "Train Epoch [6 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:26<00:00, 58.51it/s, acc=0.569, loss=1.21]\n",
            "Test  Epoch [6 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 63.42it/s, acc=0.56, loss=1.22]\n",
            "Train Epoch [7 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:26<00:00, 58.23it/s, acc=0.583, loss=1.17]\n",
            "Test  Epoch [7 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 67.69it/s, acc=0.563, loss=1.21]\n",
            "Train Epoch [8 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:27<00:00, 55.91it/s, acc=0.594, loss=1.14]\n",
            "Test  Epoch [8 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:05<00:00, 60.89it/s, acc=0.562, loss=1.22]\n",
            "Train Epoch [9 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:27<00:00, 57.10it/s, acc=0.605, loss=1.11]\n",
            "Test  Epoch [9 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 66.74it/s, acc=0.565, loss=1.22]\n",
            "Train Epoch [10 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [00:27<00:00, 56.01it/s, acc=0.611, loss=1.09]\n",
            "Test  Epoch [10 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 69.57it/s, acc=0.575, loss=1.2]\n"
          ]
        }
      ],
      "source": [
        "student = StudentNetwork(3,10).to(device)\n",
        "optimizer = torch.optim.Adam(student.parameters(), lr=lr_net)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "\n",
        "for epoch in range(student_epochs):\n",
        "\n",
        "    train_loss = AverageMeter()\n",
        "    accuracy_train = AverageMeter()\n",
        "    data_loop_train = tqdm(enumerate(trainloader), total=len(trainloader), colour='red')\n",
        "\n",
        "    student.train(True)\n",
        "    for _, train_data in data_loop_train:\n",
        "        train_img, train_label = train_data\n",
        "        train_img = train_img.to(device)\n",
        "        train_label = train_label.to(device)\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        train_pred, _ = student(train_img)\n",
        "        loss = criterion(train_pred, train_label)\n",
        "        acc = accuracy(train_pred, train_label)\n",
        "\n",
        "        train_loss.update(loss.item(), train_img.size(0))\n",
        "        accuracy_train.update(acc.item(), train_img.size(0))\n",
        "\n",
        "        dict_metrics = dict(loss = train_loss.avg, acc = accuracy_train.avg)\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        data_loop_train.set_description(f'Train Epoch [{epoch + 1} / {student_epochs}]')\n",
        "        data_loop_train.set_postfix(**dict_metrics)\n",
        "\n",
        "        for key, value in dict_metrics.items():\n",
        "              writer_student.add_scalar(f'train_{key}', value, epoch)\n",
        "\n",
        "\n",
        "    # Evaluation phase\n",
        "    student.eval()\n",
        "    data_loop_test = tqdm(enumerate(testloader), total=len(testloader),colour='green')\n",
        "    with torch.no_grad():\n",
        "\n",
        "        test_loss = AverageMeter()\n",
        "        accuracy_test = AverageMeter()\n",
        "\n",
        "        for _, test_data in data_loop_test:\n",
        "            test_img, test_label = test_data\n",
        "            test_img = test_img.to(device)\n",
        "            test_label = test_label.to(device)\n",
        "\n",
        "            test_pred, _ = student(test_img)\n",
        "            loss = criterion(test_pred, test_label)\n",
        "            acc = accuracy(test_pred, test_label)\n",
        "\n",
        "            test_loss.update(loss.item(), test_img.size(0))\n",
        "            accuracy_test.update(acc.item(), test_img.size(0))\n",
        "            dict_metrics = dict(loss = test_loss.avg, acc = accuracy_test.avg)\n",
        "\n",
        "            data_loop_test.set_description(f'Test  Epoch [{epoch + 1} / {student_epochs}]')\n",
        "            data_loop_test.set_postfix(**dict_metrics)\n",
        "\n",
        "            for key, value in dict_metrics.items():\n",
        "                writer_student.add_scalar(f'test_{key}', value, epoch)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aWlpeq7VuCCf"
      },
      "source": [
        "# Distillation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "EPPSdt2-atlx",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "19f123d8-a55a-4f99-b9ef-4c139aa6fb19"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch.autograd.anomaly_mode.set_detect_anomaly at 0x7f1d5b612fb0>"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ],
      "source": [
        "real_label = 1\n",
        "fake_label = 0\n",
        "student_distilled = StudentNetwork(3,10).to(device)\n",
        "teacher.eval()\n",
        "discriminator = Discriminator(84).to(device)\n",
        "bce_loss = nn.BCELoss()\n",
        "mse_loss = nn.MSELoss()\n",
        "optimizer_discriminator = torch.optim.Adam(discriminator.parameters(), lr=lr_net)\n",
        "optimizer_student = torch.optim.Adam(student_distilled.parameters(), lr=lr_net)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "torch.autograd.set_detect_anomaly(True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "Ed1aQQQeTaLa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c93667b6-745f-426e-de7a-bf635167c740"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Train  Epoch [1 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:48<00:00,  9.28it/s, disc_loss=1.91, stu_loss=24.9]\n",
            "Test  Epoch [1 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 65.16it/s, acc=0.395, loss=2.81]\n",
            "Train  Epoch [2 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:54<00:00,  8.94it/s, disc_loss=1.91, stu_loss=19.3]\n",
            "Test  Epoch [2 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 70.50it/s, acc=0.473, loss=2.23]\n",
            "Train  Epoch [3 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:47<00:00,  9.32it/s, disc_loss=1.91, stu_loss=17.6]\n",
            "Test  Epoch [3 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 62.85it/s, acc=0.491, loss=2.5]\n",
            "Train  Epoch [4 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:49<00:00,  9.22it/s, disc_loss=1.91, stu_loss=16.3]\n",
            "Test  Epoch [4 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 65.77it/s, acc=0.521, loss=2.28]\n",
            "Train  Epoch [5 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:48<00:00,  9.28it/s, disc_loss=1.91, stu_loss=15.3]\n",
            "Test  Epoch [5 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 62.68it/s, acc=0.548, loss=2.11]\n",
            "Train  Epoch [6 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:47<00:00,  9.33it/s, disc_loss=1.91, stu_loss=14.5]\n",
            "Test  Epoch [6 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 62.86it/s, acc=0.545, loss=2.39]\n",
            "Train  Epoch [7 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:48<00:00,  9.29it/s, disc_loss=1.91, stu_loss=13.9]\n",
            "Test  Epoch [7 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 63.95it/s, acc=0.561, loss=2.16]\n",
            "Train  Epoch [8 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:48<00:00,  9.25it/s, disc_loss=1.91, stu_loss=13.3]\n",
            "Test  Epoch [8 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 71.04it/s, acc=0.56, loss=2.29]\n",
            "Train  Epoch [9 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:49<00:00,  9.24it/s, disc_loss=1.91, stu_loss=12.9]\n",
            "Test  Epoch [9 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 71.56it/s, acc=0.565, loss=2.3]\n",
            "Train  Epoch [10 / 10]: 100%|\u001b[31m██████████\u001b[0m| 1563/1563 [02:49<00:00,  9.22it/s, disc_loss=1.91, stu_loss=12.4]\n",
            "Test  Epoch [10 / 10]: 100%|\u001b[32m██████████\u001b[0m| 313/313 [00:04<00:00, 71.51it/s, acc=0.59, loss=1.94]\n"
          ]
        }
      ],
      "source": [
        "for epoch in range(student_epochs):\n",
        "    # For each batch in the dataloader\n",
        "    data_loop_train = tqdm(enumerate(trainloader), total=len(trainloader), colour='red')\n",
        "    student_distilled.train(True)\n",
        "    train_disc_loss = AverageMeter()\n",
        "    train_stu_loss = AverageMeter()\n",
        "    for _, train_data in data_loop_train:\n",
        "\n",
        "        # UPDATE DISCRIMINATOR ADVERSARIAL\n",
        "\n",
        "        train_img, _ = train_data\n",
        "        train_img = train_img.to(device)\n",
        "        batch_size = train_img.size(0)\n",
        "        optimizer_discriminator.zero_grad()\n",
        "\n",
        "        disc_label_real = torch.full((batch_size,), real_label, dtype=torch.float).unsqueeze(-1).to(device)\n",
        "        disc_label_fake = torch.full((batch_size,), fake_label, dtype=torch.float).unsqueeze(-1).to(device)\n",
        "\n",
        "        _, features_teacher = teacher(train_img)\n",
        "        output_real = discriminator(features_teacher, False)\n",
        "        err_real = bce_loss(output_real, disc_label_real)\n",
        "\n",
        "        _, features_student = student_distilled(train_img)\n",
        "        output_fake = discriminator(features_student, False)\n",
        "        err_real = bce_loss(output_fake, disc_label_fake)\n",
        "\n",
        "        # REGULARIZER DISCRIMINATOR\n",
        "        reg = bce_loss(output_fake, disc_label_real)\n",
        "        loss = err_real + err_real + reg\n",
        "        loss.backward()\n",
        "        optimizer_discriminator.step()\n",
        "\n",
        "        # UPDATE STUDENT ADVERSARIAL + MSE\n",
        "        optimizer_student.zero_grad()\n",
        "        stu_labels = torch.full((batch_size,), real_label, dtype=torch.float).unsqueeze(-1).to(device)\n",
        "        logits_student, features_student = student_distilled(train_img)\n",
        "        logits_teacher, features_teacher = teacher(train_img)\n",
        "        output_fake = discriminator(features_student, True)\n",
        "        s_loss = bce_loss(output_fake, stu_labels) + mse_loss(logits_student, logits_teacher)\n",
        "\n",
        "        s_loss.backward()\n",
        "        optimizer_student.step()\n",
        "\n",
        "        # METRICS\n",
        "        train_disc_loss.update(loss.item(), train_img.size(0))\n",
        "        train_stu_loss.update(s_loss.item(), train_img.size(0))\n",
        "        dict_metrics = dict(disc_loss = train_disc_loss.avg, stu_loss = train_stu_loss.avg)\n",
        "\n",
        "        data_loop_train.set_description(f'Train  Epoch [{epoch + 1} / {student_epochs}]')\n",
        "        data_loop_train.set_postfix(**dict_metrics)\n",
        "\n",
        "        for key, value in dict_metrics.items():\n",
        "            writer_student_distilled.add_scalar(f'ttrain_{key}', value, epoch)\n",
        "\n",
        "\n",
        "    # Eval model\n",
        "    student_distilled.eval()\n",
        "    data_loop_test = tqdm(enumerate(testloader), total=len(testloader),colour='green')\n",
        "    with torch.no_grad():\n",
        "\n",
        "        test_loss = AverageMeter()\n",
        "        accuracy_test = AverageMeter()\n",
        "\n",
        "        for _, test_data in data_loop_test:\n",
        "            test_img, test_label = test_data\n",
        "            test_img = test_img.to(device)\n",
        "            test_label = test_label.to(device)\n",
        "\n",
        "            test_pred, _ = student_distilled(test_img)\n",
        "            loss = criterion(test_pred, test_label)\n",
        "            acc = accuracy(test_pred, test_label)\n",
        "\n",
        "            test_loss.update(loss.item(), test_img.size(0))\n",
        "            accuracy_test.update(acc.item(), test_img.size(0))\n",
        "            dict_metrics = dict(loss = test_loss.avg, acc = accuracy_test.avg)\n",
        "\n",
        "            data_loop_test.set_description(f'Test  Epoch [{epoch + 1} / {student_epochs}]')\n",
        "            data_loop_test.set_postfix(**dict_metrics)\n",
        "\n",
        "            for key, value in dict_metrics.items():\n",
        "                writer_student_distilled.add_scalar(f'test_{key}', value, epoch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "USOu27olfrsH"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "MSMyY_I2ohlO"
      },
      "execution_count": 17,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}