demo/ma/distiller/Model_Adapter_Distiller_Walkthrough_VIT_to_ResNet18_CIFAR100.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "17b6149c",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/intel/e2eAIOK/blob/main/demo/ma/distiller/Model_Adapter_Distiller_Walkthrough_VIT_to_ResNet18_CIFAR100.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b92539e",
   "metadata": {
    "id": "2b92539e"
   },
   "source": [
    "# Model Adapter Distiller Walkthrough DEMO\n",
    "Model Adapter is a convenient framework can be used to reduce training and inference time, or data labeling cost by efficiently utilizing public advanced models and those datasets from many domains. It mainly contains three components served for different cases: Finetuner, Distiller, and Domain Adapter. \n",
    "\n",
    "This demo mainly introduces the usage of Distiller. Take image classification as an example, it shows how to integrate distiller from VIT to ResNet18 on CIFAR100 dataset. This demo shows how to integrate distiller into a general training pipeline, you can find build-in simplied demo at [here](./Model_Adapter_Distiller_builtin_VIT_to_ResNet18_CIFAR100.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96ef213f",
   "metadata": {
    "id": "96ef213f"
   },
   "source": [
    "# Content\n",
    "\n",
    "* [Overview](#Overview)\n",
    "    * [Model Adapter Distiller Overview](#Model-Adapter-Distiller-Overview)\n",
    "* [Getting Started](#Getting-Started)\n",
    "    * [1. Environment Setup](#1.-Environment-Setup)\n",
    "    * [2. Data Prepare](#2.-Data-Prepare)\n",
    "    * [3. Model Prepare](#3.-Model-Prepare)\n",
    "    * [4. Train](#4.-Train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b6cb09",
   "metadata": {
    "id": "b7b6cb09"
   },
   "source": [
    "# Overview\n",
    "\n",
    "## Model Adapter Distiller Overview\n",
    "Distiller is based on knowledge distillation technology, it can transfer knowledge from a heavy model (teacher) to a light one (student) with different structure. Teacher is a large model pretrained on specific dataset, which contains sufficient knowledge for this task, while the student model has much smaller structure. Distiller trains the student not only on the dataset, but also with the help of teacher’s knowledge. With distiller, we can take use of the knowledge from the existing pretrained large models but use much less training time. It can also significantly improve the converge speed and predicting accuracy of a small model, which is very helpful for inference.\n",
    "\n",
    "<img src=\"../imgs/distiller.png\" width=\"60%\">\n",
    "<center>Model Adapter Distiller Structure</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33b8e6ac",
   "metadata": {},
   "source": [
    "# Getting Started"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "040f91d3",
   "metadata": {
    "id": "040f91d3"
   },
   "source": [
    "## 1. Environment Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d2ca573",
   "metadata": {
    "id": "9d2ca573"
   },
   "source": [
    "### (Option 1) Use Pip install\n",
    "We can directly install ModelAdapter module from Intel® End-to-End AI Optimization Kit with following command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b39bd4c",
   "metadata": {
    "id": "7b39bd4c",
    "outputId": "545c42d2-d523-4a57-fa29-820197de49bb"
   },
   "outputs": [],
   "source": [
    "!pip install e2eAIOK-ModelAdapter --pre"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f723ffa",
   "metadata": {
    "id": "3f723ffa"
   },
   "source": [
    "### (Option 2) Use Docker \n",
    "\n",
    "We can also use Docker, which contains a complete environment.\n",
    "\n",
    "Step1. prepare code\n",
    "   ``` bash\n",
    "   git clone https://github.com/intel/e2eAIOK.git\n",
    "   cd e2eAIOK\n",
    "   git submodule update --init –recursive\n",
    "   ```\n",
    "    \n",
    "Step2. build docker image\n",
    "   ``` bash\n",
    "   python3 scripts/start_e2eaiok_docker.py -b pytorch112 --dataset_path ${dataset_path} -w ${host0} ${host1} ${host2} ${host3} --proxy  \"http://addr:ip\"\n",
    "   ```\n",
    "   \n",
    "Step3. run docker and start conda env\n",
    "   ``` bash\n",
    "   sshpass -p docker ssh ${host0} -p 12347\n",
    "   conda activate pytorch-1.12.0\n",
    "   ```\n",
    "  \n",
    "Step4. Start the jupyter notebook and tensorboard service\n",
    "   ``` bash\n",
    "   nohup jupyter notebook --notebook-dir=/home/vmagent/app/e2eaiok --ip=${hostname} --port=8899 --allow-root &\n",
    "   nohup tensorboard --logdir /home/vmagent/app/data/tensorboard --host=${hostname} --port=6006 & \n",
    "   ```\n",
    "   Now you can visit demso in `http://${hostname}:8899/`, and see tensorboad log in ` http://${hostname}:6006`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec28cbf4",
   "metadata": {
    "id": "ec28cbf4"
   },
   "source": [
    "## 2. Data Prepare"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e6d8e21",
   "metadata": {
    "id": "0e6d8e21"
   },
   "source": [
    "Let's import some required modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f878badd",
   "metadata": {
    "id": "f878badd"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import transforms,datasets\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.optim as optim\n",
    "from timm.utils import accuracy\n",
    "import timm\n",
    "import transformers\n",
    "import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3893b32d",
   "metadata": {
    "id": "3893b32d"
   },
   "source": [
    "First let's define transformer for dataset, which will be needed to augment input image. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c016312d",
   "metadata": {
    "id": "c016312d"
   },
   "outputs": [],
   "source": [
    "IMAGE_MEAN = [0.5, 0.5, 0.5]\n",
    "IMAGE_STD = [0.5, 0.5, 0.5]\n",
    "\n",
    "train_transform = transforms.Compose([\n",
    "  transforms.RandomCrop(32, padding=4),\n",
    "  transforms.RandomHorizontalFlip(),\n",
    "  transforms.Resize(224),  # pretrained model is trained on large imgage size, scale 32x32 to 224x224\n",
    "  transforms.ToTensor(),\n",
    "  transforms.Normalize(IMAGE_MEAN, IMAGE_STD)\n",
    "])\n",
    "\n",
    "test_transform = transforms.Compose([\n",
    "  transforms.RandomCrop(32, padding=4),\n",
    "  transforms.Resize(224),  # pretrained model is trained on large imgage size, scale 32x32 to 224x224\n",
    "  transforms.ToTensor(),\n",
    "  transforms.Normalize(IMAGE_MEAN, IMAGE_STD)\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7422d3df",
   "metadata": {},
   "source": [
    "Then let's define CIFAR100 dataset and download it with torchvision lib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "039a91b7",
   "metadata": {
    "id": "039a91b7",
    "outputId": "b16a65c8-1512-4a9f-925b-861539a854ca"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "data_folder='./data' # dataset location\n",
    "train_set = datasets.CIFAR100(root=data_folder, train=True, download=True, transform=train_transform)\n",
    "test_set = datasets.CIFAR100(root=data_folder, train=False, download=True, transform=test_transform)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d85fb033",
   "metadata": {
    "id": "d85fb033"
   },
   "source": [
    "Finally we define dataloader, you can change batch_size and num_workers according to your own environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69896736",
   "metadata": {
    "id": "69896736"
   },
   "outputs": [],
   "source": [
    "train_loader = DataLoader(dataset=train_set, batch_size=128, shuffle=True, num_workers=1, drop_last=False)\n",
    "validate_loader = DataLoader(dataset=test_set, batch_size=128, shuffle=True, num_workers=1, drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9eeba42",
   "metadata": {
    "id": "f9eeba42"
   },
   "source": [
    "## 3. Model Prepare"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b06ea401",
   "metadata": {
    "id": "b06ea401"
   },
   "source": [
    "**Prepare Student Model**\n",
    "\n",
    "First we create a ResNe18 without pretrained weights as backbone model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40b8a7d6",
   "metadata": {
    "id": "40b8a7d6"
   },
   "outputs": [],
   "source": [
    "backbone = timm.create_model('resnet18', pretrained=False, num_classes=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1568a3ba",
   "metadata": {
    "id": "1568a3ba"
   },
   "source": [
    "(optional & recommend) Optimized weight initilization, can enhance initial learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87143657",
   "metadata": {
    "id": "87143657",
    "outputId": "3acc1696-9ef9-463e-d05d-b227b2ad413b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act1): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (drop_block): Identity()\n",
       "      (act1): ReLU(inplace=True)\n",
       "      (aa): Identity()\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act2): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=Flatten(start_dim=1, end_dim=-1))\n",
       "  (fc): Linear(in_features=512, out_features=100, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from e2eAIOK.common.trainer.model.model_utils.model_utils import initWeights\n",
    "backbone.apply(initWeights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f01b6c52",
   "metadata": {
    "id": "f01b6c52"
   },
   "source": [
    "**Prepare teacher model**\n",
    "\n",
    "To use distiller, we need to prepare teacher model to guide the training. Here we select pretrained [vit_base-224-in21k-ft-cifar100 from HuggingFace](https://huggingface.co/edumunozsala/vit_base-224-in21k-ft-cifar100)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b3d1996",
   "metadata": {
    "id": "8b3d1996"
   },
   "outputs": [],
   "source": [
    "from transformers import ViTForImageClassification\n",
    "teacher_model = ViTForImageClassification.from_pretrained('edumunozsala/vit_base-224-in21k-ft-cifar100')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6353b6d3",
   "metadata": {
    "id": "6353b6d3"
   },
   "source": [
    "**Define Distiller**\n",
    "\n",
    "Here we define a distiller using KD algorithm, and it take a teacher model as input. If teacher comes from Hugginface, please clarify \"teacher_type\" with a name starting with \"huggingface\", otherwise no need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c86c941",
   "metadata": {
    "id": "2c86c941"
   },
   "outputs": [],
   "source": [
    "from e2eAIOK.ModelAdapter.engine_core.distiller import KD\n",
    "distiller= KD(teacher_model,teacher_type=\"huggingface_vit\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5c1bd41",
   "metadata": {
    "id": "d5c1bd41"
   },
   "source": [
    "**Wrap model with Model Adapter**\n",
    "\n",
    "Call the *make_transferrable_with_knowledge_distillation()* function, which take backbone model, distiller and a loss function as input. The output model will have the ability to do the knowledge distillation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e1b2074",
   "metadata": {
    "id": "2e1b2074"
   },
   "outputs": [],
   "source": [
    "from e2eAIOK.ModelAdapter.engine_core.transferrable_model import *\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "model = make_transferrable_with_knowledge_distillation(backbone,loss_fn,distiller)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e633199",
   "metadata": {
    "id": "4e633199"
   },
   "source": [
    "## 4. Train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa226f93",
   "metadata": {
    "id": "aa226f93"
   },
   "source": [
    "**Create optimizer**\n",
    "\n",
    "Create the Optimizer, here we choose SGD."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c9561a5",
   "metadata": {
    "id": "5c9561a5"
   },
   "outputs": [],
   "source": [
    "init_lr = 0.01\n",
    "weight_decay = 0.0001\n",
    "momentum = 0.9\n",
    "optimizer = optim.SGD(model.parameters(),lr=init_lr, weight_decay=weight_decay,momentum=momentum)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e23ccfe",
   "metadata": {},
   "source": [
    "**Create scheduler**\n",
    "\n",
    "Here we choose a *ExponentialLR* scheduler, you can also change to other schedulers for your own task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db8e784a",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "840a0628",
   "metadata": {
    "id": "840a0628"
   },
   "source": [
    "**Create Trainer**\n",
    "\n",
    "Create a simple *Trainer*, which contains *train()* and *evaluate()* function for this simple ResNet50 training task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "874fb6b2",
   "metadata": {
    "id": "874fb6b2"
   },
   "outputs": [],
   "source": [
    "max_epoch = 1 # max 1 epoch\n",
    "print_interval = 10 \n",
    "\n",
    "class Trainer:\n",
    "    def __init__(self, model, optimizer, scheduler):\n",
    "        self._model = model\n",
    "        self._optimizer = optimizer\n",
    "        self._scheduler = scheduler\n",
    "        \n",
    "    def train(self, train_dataloader, valid_dataloader, max_epoch):\n",
    "        ''' \n",
    "        :param train_dataloader: train dataloader\n",
    "        :param valid_dataloader: validation dataloader\n",
    "        :param max_epoch: steps per epoch\n",
    "        '''\n",
    "        for epoch in range(0, max_epoch):\n",
    "            ################## train #####################\n",
    "            self._model.train()  # set training flag\n",
    "            for (cur_step,(data, label)) in enumerate(train_dataloader):\n",
    "                self._optimizer.zero_grad()\n",
    "                output = self._model(data)\n",
    "                loss_value = self._model.loss(output, label) # transferrable model has loss attribute\n",
    "                loss_value.backward() \n",
    "                if cur_step%print_interval == 0:\n",
    "                    batch_acc = accuracy(output.backbone_output,label)[0] # use output.backbone_output instead of output\n",
    "                    dt = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') # date time\n",
    "                    print(\"[{}] epoch {} step {} : total loss {:4f}, training backbone loss {:.4f}, distiller loss {:4f}, training batch acc {:.4f}\".format(\n",
    "                      dt, epoch, cur_step, loss_value.total_loss.item(),loss_value.backbone_loss.item(), loss_value.distiller_loss.item(), batch_acc.item())) \n",
    "                self._optimizer.step()\n",
    "            self._scheduler.step()\n",
    "            ################## evaluate ######################\n",
    "            self.evaluate(valid_dataloader)\n",
    "            \n",
    "    def evaluate(self, valid_dataloader):\n",
    "        with torch.no_grad():\n",
    "            self._model.eval()  \n",
    "            backbone = self._model.backbone # use backbone in evaluation\n",
    "            loss_cum = 0.0\n",
    "            sample_num = 0\n",
    "            acc_cum = 0.0\n",
    "            total_step = len(valid_dataloader)\n",
    "            for (cur_step,(data, label)) in enumerate(valid_dataloader):\n",
    "                output = backbone(data)\n",
    "                batch_size = data.size(0)\n",
    "                sample_num += batch_size\n",
    "                loss_cum += loss_fn(output, label).item() * batch_size\n",
    "                acc_cum += accuracy(output, label)[0].item() * batch_size\n",
    "                if cur_step%print_interval == 0:\n",
    "                    print(f\"step {cur_step}/{total_step}\")\n",
    "            dt = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') # date time\n",
    "            loss_value = loss_cum/sample_num\n",
    "            acc_value = acc_cum/sample_num\n",
    "\n",
    "            print(\"[{}] evaluation loss {:.4f}, evaluation acc {:.4f}\".format(\n",
    "                dt, loss_value, acc_value))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc734272",
   "metadata": {
    "id": "bc734272"
   },
   "source": [
    "**Train and Evaluate**\n",
    "\n",
    "Let's start trainer, train for 1 epoch and evaluate the final accuracy.\n",
    "\n",
    "This will take some time, have a break and get a coffee!\n",
    "\n",
    "Here we only train one epoch for a quick test, you may expect a result with accuracy around 0.15~0.17\n",
    "\n",
    "You can get an optimized and accelerated training with saving logits function, refer to [logits saving demo](Model_Adapter_Distiller_customized_ResNet18_CIFAR100_train_with_logits.ipynb) and [training with saved logits demo](./Model_Adapter_Distiller_customized_ResNet18_CIFAR100_train_with_logits.ipynb) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bae383b6",
   "metadata": {
    "id": "bae383b6",
    "outputId": "24b77f05-5af2-43c5-bfd7-b510868e36c9",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-02-13 05:49:36] epoch 0 step 0 : total loss 3.099823, training backbone loss 5.4708, distiller loss 2.836383, training batch acc 0.0000\n",
      "[2023-02-13 05:50:46] epoch 0 step 10 : total loss 2.362927, training backbone loss 4.5393, distiller loss 2.121112, training batch acc 1.5625\n",
      "[2023-02-13 05:51:57] epoch 0 step 20 : total loss 2.394127, training backbone loss 4.4548, distiller loss 2.165160, training batch acc 2.3438\n",
      "[2023-02-13 05:53:06] epoch 0 step 30 : total loss 2.356652, training backbone loss 4.4124, distiller loss 2.128232, training batch acc 3.9062\n",
      "[2023-02-13 05:54:14] epoch 0 step 40 : total loss 2.320713, training backbone loss 4.4267, distiller loss 2.086718, training batch acc 4.6875\n",
      "[2023-02-13 05:55:21] epoch 0 step 50 : total loss 2.227285, training backbone loss 4.2046, distiller loss 2.007580, training batch acc 6.2500\n",
      "[2023-02-13 05:56:33] epoch 0 step 60 : total loss 2.237528, training backbone loss 4.3497, distiller loss 2.002843, training batch acc 5.4688\n",
      "[2023-02-13 05:57:44] epoch 0 step 70 : total loss 2.239463, training backbone loss 4.2061, distiller loss 2.020943, training batch acc 8.5938\n",
      "[2023-02-13 05:58:55] epoch 0 step 80 : total loss 2.262591, training backbone loss 4.2680, distiller loss 2.039762, training batch acc 7.8125\n",
      "[2023-02-13 06:00:05] epoch 0 step 90 : total loss 2.337352, training backbone loss 4.3197, distiller loss 2.117093, training batch acc 3.9062\n",
      "[2023-02-13 06:01:16] epoch 0 step 100 : total loss 2.195452, training backbone loss 4.0890, distiller loss 1.985061, training batch acc 10.1562\n",
      "[2023-02-13 06:02:27] epoch 0 step 110 : total loss 2.126554, training backbone loss 3.9784, distiller loss 1.920798, training batch acc 10.1562\n",
      "[2023-02-13 06:03:39] epoch 0 step 120 : total loss 2.153989, training backbone loss 3.9720, distiller loss 1.951988, training batch acc 12.5000\n",
      "[2023-02-13 06:04:50] epoch 0 step 130 : total loss 2.223530, training backbone loss 4.0605, distiller loss 2.019420, training batch acc 10.1562\n",
      "[2023-02-13 06:05:57] epoch 0 step 140 : total loss 2.178571, training backbone loss 3.9980, distiller loss 1.976413, training batch acc 12.5000\n",
      "[2023-02-13 06:07:06] epoch 0 step 150 : total loss 2.163835, training backbone loss 3.9886, distiller loss 1.961078, training batch acc 16.4062\n",
      "[2023-02-13 06:08:16] epoch 0 step 160 : total loss 2.137477, training backbone loss 3.8332, distiller loss 1.949065, training batch acc 11.7188\n",
      "[2023-02-13 06:09:24] epoch 0 step 170 : total loss 2.100381, training backbone loss 3.8991, distiller loss 1.900521, training batch acc 13.2812\n",
      "[2023-02-13 06:10:34] epoch 0 step 180 : total loss 2.214239, training backbone loss 4.0391, distiller loss 2.011480, training batch acc 6.2500\n",
      "[2023-02-13 06:11:45] epoch 0 step 190 : total loss 2.126791, training backbone loss 3.9302, distiller loss 1.926417, training batch acc 14.0625\n",
      "[2023-02-13 06:12:57] epoch 0 step 200 : total loss 2.076819, training backbone loss 3.8596, distiller loss 1.878730, training batch acc 15.6250\n",
      "[2023-02-13 06:14:06] epoch 0 step 210 : total loss 2.229523, training backbone loss 4.0190, distiller loss 2.030696, training batch acc 11.7188\n",
      "[2023-02-13 06:15:15] epoch 0 step 220 : total loss 2.130007, training backbone loss 3.9382, distiller loss 1.929101, training batch acc 10.9375\n",
      "[2023-02-13 06:16:25] epoch 0 step 230 : total loss 2.197126, training backbone loss 3.9229, distiller loss 2.005373, training batch acc 8.5938\n",
      "[2023-02-13 06:17:36] epoch 0 step 240 : total loss 2.134550, training backbone loss 3.8970, distiller loss 1.938721, training batch acc 14.0625\n",
      "[2023-02-13 06:18:47] epoch 0 step 250 : total loss 2.115740, training backbone loss 3.7911, distiller loss 1.929584, training batch acc 14.0625\n",
      "[2023-02-13 06:19:59] epoch 0 step 260 : total loss 2.162165, training backbone loss 3.9435, distiller loss 1.964244, training batch acc 13.2812\n",
      "[2023-02-13 06:21:08] epoch 0 step 270 : total loss 2.195049, training backbone loss 4.0198, distiller loss 1.992294, training batch acc 14.0625\n",
      "[2023-02-13 06:22:18] epoch 0 step 280 : total loss 2.129143, training backbone loss 3.8038, distiller loss 1.943075, training batch acc 13.2812\n",
      "[2023-02-13 06:23:27] epoch 0 step 290 : total loss 2.056203, training backbone loss 3.7442, distiller loss 1.868651, training batch acc 14.0625\n",
      "[2023-02-13 06:24:35] epoch 0 step 300 : total loss 2.053657, training backbone loss 3.8061, distiller loss 1.858943, training batch acc 13.2812\n",
      "[2023-02-13 06:25:44] epoch 0 step 310 : total loss 2.043147, training backbone loss 3.7814, distiller loss 1.850006, training batch acc 9.3750\n",
      "[2023-02-13 06:26:51] epoch 0 step 320 : total loss 2.167141, training backbone loss 3.7805, distiller loss 1.987880, training batch acc 11.7188\n",
      "[2023-02-13 06:27:57] epoch 0 step 330 : total loss 2.176765, training backbone loss 3.9466, distiller loss 1.980119, training batch acc 10.9375\n",
      "[2023-02-13 06:29:04] epoch 0 step 340 : total loss 2.037331, training backbone loss 3.7802, distiller loss 1.843684, training batch acc 17.1875\n",
      "[2023-02-13 06:30:11] epoch 0 step 350 : total loss 2.066255, training backbone loss 3.6719, distiller loss 1.887855, training batch acc 16.4062\n",
      "[2023-02-13 06:31:19] epoch 0 step 360 : total loss 2.081556, training backbone loss 3.7981, distiller loss 1.890828, training batch acc 10.1562\n",
      "[2023-02-13 06:32:28] epoch 0 step 370 : total loss 2.112749, training backbone loss 3.7865, distiller loss 1.926774, training batch acc 15.6250\n",
      "[2023-02-13 06:33:38] epoch 0 step 380 : total loss 2.016887, training backbone loss 3.5821, distiller loss 1.842977, training batch acc 19.5312\n",
      "[2023-02-13 06:34:44] epoch 0 step 390 : total loss 2.084308, training backbone loss 3.8419, distiller loss 1.889014, training batch acc 12.5000\n",
      "step 0/79\n",
      "step 10/79\n",
      "step 20/79\n",
      "step 30/79\n",
      "step 40/79\n",
      "step 50/79\n",
      "step 60/79\n",
      "step 70/79\n",
      "[2023-02-13 06:35:21] evaluation loss 3.7059, evaluation acc 16.0000\n",
      "CPU times: user 22h 28min 7s, sys: 7h 33min 26s, total: 1d 6h 1min 33s\n",
      "Wall time: 45min 52s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "trainer = Trainer(model, optimizer, scheduler)\n",
    "trainer.train(train_loader,validate_loader,max_epoch)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "provenance": []
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}