hw7_Network_Compression/hw7_Knowledge_Distillation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "b5cFq_TgWlQ_"
   },
   "source": [
    "# Homework 7 - Network Compression (Knowledge Distillation)\n",
    "\n",
    "网络压缩方法：知识蒸馏"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vNiZCGrIYKdR"
   },
   "source": [
    "# Readme\n",
    "\n",
    "Compression有很多方法，作业7实现其中的四种，分別是:\n",
    "\n",
    "* 知识蒸馏 Knowledge Distillation\n",
    "* 网络剪枝 Network Pruning\n",
    "* 结构设计 Architecture Design\n",
    "* 参数量话 Weight Quantization\n",
    "\n",
    "在這個notebook中我們會介紹Knowledge Distillation，\n",
    "而我們有提供已經學習好的大model方便大家做Knowledge Distillation。\n",
    "而我們使用的小model是\"Architecture Design\"過的model。\n",
    "\n",
    "* Model结构 / Architecute Design在同目录的hw7_Architecture_Design.ipynb。\n",
    "* 下載pretrained大model(47.2M): https://drive.google.com/uc?id=1B8ljdrxYXJsZv2vmTequdPOofp3VF3NN\n",
    "  * 請使用torchvision提供的ResNet18，把num_classes改成11後load進去即可。(後面有範例。)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 84
    },
    "colab_type": "code",
    "id": "XdzskhdEb65Z",
    "outputId": "e951b5f4-f5f0-4e36-83f3-2997a4eca14b"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models\n",
    "# Load進我們的Model架構(在hw7_Architecture_Design.ipynb內)\n",
    "%run \"hw7_Architecture_Design.ipynb\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bdUtCxBBcH0B"
   },
   "source": [
    "Knowledge Distillation\n",
    "===\n",
    "\n",
    "<img src=\"https://i.imgur.com/H2aF7Rv.png=100x\" width=\"500px\">\n",
    "\n",
    "簡單上來說就是讓已經做得很好的大model們去告訴小model\"如何\"學習。\n",
    "而我們如何做到這件事情呢? 就是利用大model預測的logits給小model當作標準就可以了。\n",
    "\n",
    "## 為甚麼這會work?\n",
    "* 例如當data不是很乾淨的時候，對一般的model來說他是個noise，只會干擾學習。透過去學習其他大model預測的logits會比較好。\n",
    "* label和label之間可能有關連，這可以引導小model去學習。例如數字8可能就和6,9,0有關係。\n",
    "* 弱化已經學習不錯的target(?)，避免讓其gradient干擾其他還沒學好的task。\n",
    "\n",
    "\n",
    "## 要怎麼實作?\n",
    "* $Loss = \\alpha T^2 \\times KL(\\frac{\\text{Teacher's Logits}}{T} || \\frac{\\text{Student's Logits}}{T}) + (1-\\alpha) \\times CE(\\text{ labels } || \\text{ Student's Logits })$\n",
    "\n",
    "\n",
    "* 以下code為甚麼要對student使用log_softmax: https://github.com/peterliht/knowledge-distillation-pytorch/issues/2\n",
    "* reference: [Distilling the Knowledge in a Neural Network](https://arxiv.org/abs/1503.02531)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "M-dSi_P-4les"
   },
   "outputs": [],
   "source": [
    "def loss_fn_kd(outputs, labels, teacher_outputs, T=20, alpha=0.5):\n",
    "    # 一般的Cross Entropy\n",
    "    hard_loss = F.cross_entropy(outputs, labels) * (1. - alpha)\n",
    "    # 讓logits的log_softmax對目標機率(teacher的logits/T後softmax)做KL Divergence。\n",
    "    soft_loss = nn.KLDivLoss(reduction='batchmean')(F.log_softmax(outputs/T, dim=1),\n",
    "                             F.softmax(teacher_outputs/T, dim=1)) * (alpha * T * T)\n",
    "    return hard_loss + soft_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NfnRoOt5VIze"
   },
   "source": [
    "# Data Processing\n",
    "\n",
    "我們的Dataset使用的是跟Hw3 - CNN同樣的Dataset，因此這個區塊的Augmentation / Read Image大家參考或直接抄就好。\n",
    "\n",
    "如果有不會的話可以回去看Hw3的colab。\n",
    "\n",
    "需要注意的是如果要自己寫的話，Augment的方法最好使用我們的方法，避免輸入有差異導致Teacher Net預測不好。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ExdUvTRaVNOT"
   },
   "outputs": [],
   "source": [
    "import re\n",
    "import torch\n",
    "from glob import glob\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "class MyDataset(torch.utils.data.Dataset):\n",
    "\n",
    "    def __init__(self, folderName, transform=None):\n",
    "        self.transform = transform\n",
    "        self.data = []\n",
    "        self.label = []\n",
    "\n",
    "        for img_path in glob(folderName + '/*.jpg'):\n",
    "            try:\n",
    "                # Get classIdx by parsing image path\n",
    "                class_idx = int(re.findall(re.compile(r'\\d+'), img_path)[1])\n",
    "            except:\n",
    "                # if inference mode (there's no answer), class_idx default 0\n",
    "                class_idx = 0\n",
    " \n",
    "            image = Image.open(img_path)\n",
    "            # Get File Descriptor\n",
    "            image_fp = image.fp\n",
    "            image.load()\n",
    "            # Close File Descriptor (or it'll reach OPEN_MAX)\n",
    "            image_fp.close()\n",
    "\n",
    "            self.data.append(image)\n",
    "            self.label.append(class_idx)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "        image = self.data[idx]\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image, self.label[idx]\n",
    "\n",
    "\n",
    "trainTransform = transforms.Compose([\n",
    "    transforms.RandomCrop(256, pad_if_needed=True, padding_mode='symmetric'),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomRotation(15),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "testTransform = transforms.Compose([\n",
    "    transforms.CenterCrop(256),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "def get_dataloader(mode='training', batch_size=32):\n",
    "\n",
    "    assert mode in ['training', 'evaluation', 'validation']\n",
    "\n",
    "    dataset = MyDataset(\n",
    "        f'./food-11/{mode}',\n",
    "        transform=trainTransform if mode == 'training' else testTransform)\n",
    "\n",
    "    dataloader = torch.utils.data.DataLoader(\n",
    "        dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=(mode == 'training'))\n",
    "\n",
    "    return dataloader\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ACPwL9_JWceQ"
   },
   "source": [
    "# Pre-processing\n",
    "\n",
    "我們已經提供TeacherNet的state_dict，其架構是torchvision提供的ResNet18。\n",
    "\n",
    "至於StudentNet的架構則在hw7_Architecture_Design.ipynb中。\n",
    "\n",
    "這裡我們使用的Optimizer為AdamW，沒有為甚麼，就純粹我想用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wzuuGvnbWkG8"
   },
   "outputs": [],
   "source": [
    "# get dataloader\n",
    "train_dataloader = get_dataloader('training', batch_size=32)\n",
    "valid_dataloader = get_dataloader('validation', batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 84
    },
    "colab_type": "code",
    "id": "ZWdQtDtgoGCp",
    "outputId": "ed020dd3-2875-48bc-b4c6-7ab8e46e314c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=1B8ljdrxYXJsZv2vmTequdPOofp3VF3NN\n",
      "To: /content/teacher_resnet18.bin\n",
      "44.8MB [00:00, 170MB/s]\n"
     ]
    }
   ],
   "source": [
    "teacher_net = models.resnet18(pretrained=False, num_classes=11).cuda()\n",
    "student_net = StudentNet(base=16).cuda()\n",
    "\n",
    "teacher_net.load_state_dict(torch.load(f'./teacher_resnet18.bin'))\n",
    "optimizer = optim.AdamW(student_net.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Wvc1W5yO2QaE"
   },
   "source": [
    "# Start Training\n",
    "\n",
    "* 剩下的步驟與你在做Hw3 - CNN的時候一樣。\n",
    "\n",
    "## 小提醒\n",
    "\n",
    "* torch.no_grad是指接下來的運算或該tensor不需要算gradient。\n",
    "* model.eval()與model.train()差在於Batchnorm要不要紀錄，以及要不要做Dropout。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "-TzmWtT62Qmy",
    "outputId": "70ed38d7-e0e5-45b5-8dc6-3febc6e4e1b9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch   0: train loss: 15.0518, acc 0.3109 valid loss: 16.0725, acc 0.3539\n",
      "epoch   1: train loss: 13.4940, acc 0.3987 valid loss: 14.7710, acc 0.4367\n",
      "epoch   2: train loss: 12.4698, acc 0.4525 valid loss: 14.4754, acc 0.4548\n",
      "epoch   3: train loss: 11.6713, acc 0.4806 valid loss: 12.2971, acc 0.5029\n",
      "epoch   4: train loss: 11.2376, acc 0.5075 valid loss: 12.3508, acc 0.5195\n",
      "epoch   5: train loss: 10.5119, acc 0.5350 valid loss: 10.7621, acc 0.5659\n",
      "epoch   6: train loss: 10.1884, acc 0.5594 valid loss: 10.0260, acc 0.5866\n",
      "epoch   7: train loss: 9.7195, acc 0.5701 valid loss: 9.7483, acc 0.5933\n",
      "epoch   8: train loss: 9.2268, acc 0.5932 valid loss: 11.3290, acc 0.5711\n",
      "epoch   9: train loss: 8.9704, acc 0.6038 valid loss: 9.3383, acc 0.6210\n",
      "epoch  10: train loss: 8.6701, acc 0.6151 valid loss: 8.6603, acc 0.6461\n",
      "epoch  11: train loss: 8.2406, acc 0.6337 valid loss: 9.7204, acc 0.6105\n",
      "epoch  12: train loss: 8.1295, acc 0.6328 valid loss: 8.1776, acc 0.6647\n",
      "epoch  13: train loss: 8.0152, acc 0.6404 valid loss: 8.1389, acc 0.6805\n",
      "epoch  14: train loss: 7.6996, acc 0.6544 valid loss: 8.6098, acc 0.6455\n",
      "epoch  15: train loss: 7.4576, acc 0.6609 valid loss: 7.1417, acc 0.6965\n",
      "epoch  16: train loss: 7.4529, acc 0.6704 valid loss: 9.2468, acc 0.6251\n",
      "epoch  17: train loss: 7.2072, acc 0.6797 valid loss: 7.3580, acc 0.6816\n",
      "epoch  18: train loss: 7.1122, acc 0.6790 valid loss: 7.7546, acc 0.6682\n",
      "epoch  19: train loss: 7.0656, acc 0.6827 valid loss: 7.6994, acc 0.6682\n",
      "epoch  20: train loss: 6.9064, acc 0.6932 valid loss: 7.3854, acc 0.6930\n",
      "epoch  21: train loss: 6.8511, acc 0.6947 valid loss: 6.7831, acc 0.7096\n",
      "epoch  22: train loss: 6.6618, acc 0.7057 valid loss: 6.6225, acc 0.7114\n",
      "epoch  23: train loss: 6.4619, acc 0.7081 valid loss: 7.1267, acc 0.7076\n",
      "epoch  24: train loss: 6.5381, acc 0.7044 valid loss: 6.3144, acc 0.7303\n",
      "epoch  25: train loss: 6.3831, acc 0.7184 valid loss: 6.2740, acc 0.7166\n",
      "epoch  26: train loss: 6.2874, acc 0.7147 valid loss: 6.5714, acc 0.7155\n",
      "epoch  27: train loss: 6.1034, acc 0.7243 valid loss: 6.6524, acc 0.7058\n",
      "epoch  28: train loss: 6.1374, acc 0.7205 valid loss: 6.6118, acc 0.7140\n",
      "epoch  29: train loss: 6.0424, acc 0.7272 valid loss: 6.7327, acc 0.7187\n",
      "epoch  30: train loss: 5.9303, acc 0.7326 valid loss: 5.9404, acc 0.7379\n",
      "epoch  31: train loss: 5.9641, acc 0.7284 valid loss: 7.4335, acc 0.7067\n",
      "epoch  32: train loss: 5.9078, acc 0.7381 valid loss: 6.1215, acc 0.7324\n",
      "epoch  33: train loss: 5.7420, acc 0.7436 valid loss: 6.2652, acc 0.7359\n",
      "epoch  34: train loss: 5.8273, acc 0.7438 valid loss: 6.3690, acc 0.7356\n",
      "epoch  35: train loss: 5.6804, acc 0.7456 valid loss: 5.7503, acc 0.7443\n",
      "epoch  36: train loss: 5.5746, acc 0.7475 valid loss: 6.1845, acc 0.7210\n",
      "epoch  37: train loss: 5.6103, acc 0.7591 valid loss: 5.2901, acc 0.7510\n",
      "epoch  38: train loss: 5.5407, acc 0.7548 valid loss: 6.0876, acc 0.7423\n",
      "epoch  39: train loss: 5.4807, acc 0.7587 valid loss: 6.2710, acc 0.7376\n",
      "epoch  40: train loss: 5.4552, acc 0.7641 valid loss: 5.7196, acc 0.7557\n",
      "epoch  41: train loss: 5.3654, acc 0.7652 valid loss: 5.6415, acc 0.7499\n",
      "epoch  42: train loss: 5.3321, acc 0.7643 valid loss: 5.8577, acc 0.7501\n",
      "epoch  43: train loss: 5.3093, acc 0.7622 valid loss: 5.4543, acc 0.7525\n",
      "epoch  44: train loss: 5.2165, acc 0.7705 valid loss: 5.6573, acc 0.7603\n",
      "epoch  45: train loss: 5.2933, acc 0.7733 valid loss: 5.6809, acc 0.7548\n",
      "epoch  46: train loss: 5.1821, acc 0.7732 valid loss: 5.3498, acc 0.7592\n",
      "epoch  47: train loss: 5.0384, acc 0.7820 valid loss: 5.4590, acc 0.7557\n",
      "epoch  48: train loss: 5.1103, acc 0.7808 valid loss: 5.6532, acc 0.7545\n",
      "epoch  49: train loss: 5.0507, acc 0.7846 valid loss: 5.7793, acc 0.7522\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "ignored",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-936f81540f69>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m200\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m     \u001b[0mstudent_net\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 38\u001b[0;31m     \u001b[0mtrain_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_dataloader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mupdate\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     39\u001b[0m     \u001b[0mstudent_net\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m     \u001b[0mvalid_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalid_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalid_dataloader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mupdate\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-8-936f81540f69>\u001b[0m in \u001b[0;36mrun_epoch\u001b[0;34m(dataloader, update, alpha)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mrun_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataloader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mupdate\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0malpha\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m     \u001b[0mtotal_num\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtotal_hit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtotal_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m     \u001b[0;32mfor\u001b[0m \u001b[0mnow_step\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_data\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataloader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m         \u001b[0;31m# 清空 optimizer\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    344\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__next__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    345\u001b[0m         \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 346\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_fetcher\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    347\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    348\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "def run_epoch(dataloader, update=True, alpha=0.5):\n",
    "    total_num, total_hit, total_loss = 0, 0, 0\n",
    "    for now_step, batch_data in enumerate(dataloader):\n",
    "        # 清空 optimizer\n",
    "        optimizer.zero_grad()\n",
    "        # 處理 input\n",
    "        inputs, hard_labels = batch_data\n",
    "        inputs = inputs.cuda()\n",
    "        hard_labels = torch.LongTensor(hard_labels).cuda()\n",
    "        # 因為Teacher沒有要backprop，所以我們使用torch.no_grad\n",
    "        # 告訴torch不要暫存中間值(去做backprop)以浪費記憶體空間。\n",
    "        with torch.no_grad():\n",
    "            soft_labels = teacher_net(inputs)\n",
    "\n",
    "        if update:\n",
    "            logits = student_net(inputs)\n",
    "            # 使用我們之前所寫的融合soft label&hard label的loss。\n",
    "            # T=20是原始論文的參數設定。\n",
    "            loss = loss_fn_kd(logits, hard_labels, soft_labels, 20, alpha)\n",
    "            loss.backward()\n",
    "            optimizer.step()    \n",
    "        else:\n",
    "            # 只是算validation acc的話，就開no_grad節省空間。\n",
    "            with torch.no_grad():\n",
    "                logits = student_net(inputs)\n",
    "                loss = loss_fn_kd(logits, hard_labels, soft_labels, 20, alpha)\n",
    "            \n",
    "        total_hit += torch.sum(torch.argmax(logits, dim=1) == hard_labels).item()\n",
    "        total_num += len(inputs)\n",
    "\n",
    "        total_loss += loss.item() * len(inputs)\n",
    "    return total_loss / total_num, total_hit / total_num\n",
    "\n",
    "\n",
    "# TeacherNet永遠都是Eval mode.\n",
    "teacher_net.eval()\n",
    "now_best_acc = 0\n",
    "for epoch in range(200):\n",
    "    student_net.train()\n",
    "    train_loss, train_acc = run_epoch(train_dataloader, update=True)\n",
    "    student_net.eval()\n",
    "    valid_loss, valid_acc = run_epoch(valid_dataloader, update=False)\n",
    "\n",
    "    # 存下最好的model。\n",
    "    if valid_acc > now_best_acc:\n",
    "        now_best_acc = valid_acc\n",
    "        torch.save(student_net.state_dict(), 'student_model.bin')\n",
    "    print('epoch {:>3d}: train loss: {:6.4f}, acc {:6.4f} valid loss: {:6.4f}, acc {:6.4f}'.format(\n",
    "        epoch, train_loss, train_acc, valid_loss, valid_acc))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0GObCiGNtPkZ"
   },
   "source": [
    "# Inference\n",
    "\n",
    "同Hw3，請參考該作業:)。\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "hw7_Knowledge_Distillation.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}