DML_5S_only.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-20T01:22:29.368Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/data/yaliu/jupyterbooks/multi-KD/models/teacher/resnet20.py:36: UserWarning: nn.init.kaiming_normal is now deprecated in favor of nn.init.kaiming_normal_.\n",
      "  init.kaiming_normal(m.weight)\n",
      "WARNING:root:Setting up a new session...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "*-----------------DML----------------*\n",
      "\n",
      "===> epoch: 1/200\n",
      "Training:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:226: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:227: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:228: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:229: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:230: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/torch/nn/functional.py:1992: UserWarning: reduction: 'mean' divides the total loss by both the batch size and the support size.'batchmean' divides only by the batch size, and aligns with the KL div math definition.'mean' will be changed to behave the same as 'batchmean' in the next major release.\n",
      "  warnings.warn(\"reduction: 'mean' divides the total loss by both the batch size and the support size.\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/391]\tTime 0.473 (0.473)\tData 0.039 (0.039)\tLoss 2.3204 (2.3204)\tPrec@1 (7.812)\n",
      "[40/391]\tTime 0.275 (0.281)\tData 0.019 (0.020)\tLoss 1.5083 (1.6540)\tPrec@1 (22.847)\n",
      "[80/391]\tTime 0.275 (0.280)\tData 0.019 (0.019)\tLoss 1.2872 (1.5041)\tPrec@1 (28.241)\n",
      "[120/391]\tTime 0.284 (0.280)\tData 0.019 (0.019)\tLoss 1.2726 (1.4197)\tPrec@1 (31.431)\n",
      "[160/391]\tTime 0.289 (0.280)\tData 0.019 (0.019)\tLoss 1.1916 (1.3726)\tPrec@1 (33.438)\n",
      "[200/391]\tTime 0.289 (0.282)\tData 0.020 (0.020)\tLoss 1.0543 (1.3281)\tPrec@1 (35.393)\n",
      "[240/391]\tTime 0.284 (0.283)\tData 0.019 (0.019)\tLoss 1.0713 (1.2909)\tPrec@1 (37.017)\n",
      "[280/391]\tTime 0.293 (0.285)\tData 0.019 (0.020)\tLoss 1.0685 (1.2523)\tPrec@1 (38.723)\n",
      "[320/391]\tTime 0.306 (0.286)\tData 0.021 (0.020)\tLoss 0.8693 (1.2183)\tPrec@1 (40.301)\n",
      "[360/391]\tTime 0.284 (0.287)\tData 0.019 (0.020)\tLoss 0.8677 (1.1847)\tPrec@1 (41.915)\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.2756 (1.2756)\tPrec@1 53.125 (53.125)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 1.4258 (1.3573)\tPrec@1 49.219 (51.963)\n",
      " * ResNet20 Prec@1 51.290\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.7626 (1.7626)\tPrec@1 42.188 (42.188)\n",
      "Test: [40/79]\tTime 0.019 (0.019)\tLoss 1.8302 (1.8178)\tPrec@1 39.844 (41.139)\n",
      " * ResNet20 Prec@1 40.480\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.3397 (1.3397)\tPrec@1 53.125 (53.125)\n",
      "Test: [40/79]\tTime 0.019 (0.019)\tLoss 1.3498 (1.3611)\tPrec@1 52.344 (51.162)\n",
      " * ResNet20 Prec@1 51.050\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.7008 (1.7008)\tPrec@1 45.312 (45.312)\n",
      "Test: [40/79]\tTime 0.019 (0.020)\tLoss 1.8361 (1.7284)\tPrec@1 46.094 (45.141)\n",
      " * ResNet20 Prec@1 44.600\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.019 (0.019)\tLoss 1.2495 (1.2495)\tPrec@1 51.562 (51.562)\n",
      "Test: [40/79]\tTime 0.020 (0.019)\tLoss 1.1314 (1.2524)\tPrec@1 55.469 (55.030)\n",
      " * ResNet20 Prec@1 54.670\n",
      "\n",
      "===> epoch: 2/200\n",
      "Training:\n",
      "[0/391]\tTime 0.308 (0.308)\tData 0.020 (0.020)\tLoss 0.7366 (0.7366)\tPrec@1 (61.719)\n",
      "[40/391]\tTime 0.302 (0.290)\tData 0.020 (0.019)\tLoss 0.7652 (0.8396)\tPrec@1 (57.927)\n",
      "[80/391]\tTime 0.294 (0.293)\tData 0.019 (0.019)\tLoss 0.7573 (0.8381)\tPrec@1 (57.967)\n",
      "[120/391]\tTime 0.285 (0.292)\tData 0.019 (0.019)\tLoss 0.8173 (0.8184)\tPrec@1 (58.697)\n",
      "[160/391]\tTime 0.293 (0.295)\tData 0.019 (0.019)\tLoss 0.6966 (0.8004)\tPrec@1 (59.399)\n",
      "[200/391]\tTime 0.298 (0.295)\tData 0.019 (0.019)\tLoss 0.8094 (0.7870)\tPrec@1 (60.117)\n",
      "[240/391]\tTime 0.284 (0.296)\tData 0.019 (0.019)\tLoss 0.7279 (0.7768)\tPrec@1 (60.555)\n",
      "[280/391]\tTime 0.291 (0.296)\tData 0.020 (0.019)\tLoss 0.6203 (0.7600)\tPrec@1 (61.277)\n",
      "[320/391]\tTime 0.298 (0.296)\tData 0.020 (0.020)\tLoss 0.6156 (0.7548)\tPrec@1 (61.556)\n",
      "[360/391]\tTime 0.294 (0.297)\tData 0.020 (0.019)\tLoss 0.5917 (0.7452)\tPrec@1 (62.037)\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.3291 (1.3291)\tPrec@1 57.031 (57.031)\n",
      "Test: [40/79]\tTime 0.019 (0.019)\tLoss 1.2574 (1.3647)\tPrec@1 58.594 (55.488)\n",
      " * ResNet20 Prec@1 55.230\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.2256 (1.2256)\tPrec@1 57.031 (57.031)\n",
      "Test: [40/79]\tTime 0.019 (0.020)\tLoss 1.2151 (1.2919)\tPrec@1 55.469 (55.069)\n",
      " * ResNet20 Prec@1 54.370\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.0608 (1.0608)\tPrec@1 58.594 (58.594)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 1.1017 (1.1839)\tPrec@1 61.719 (57.470)\n",
      " * ResNet20 Prec@1 58.040\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.019 (0.019)\tLoss 1.0091 (1.0091)\tPrec@1 64.062 (64.062)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 0.9900 (1.0569)\tPrec@1 60.156 (62.405)\n",
      " * ResNet20 Prec@1 62.270\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.019 (0.019)\tLoss 0.9638 (0.9638)\tPrec@1 66.406 (66.406)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 1.0202 (1.0499)\tPrec@1 65.625 (62.557)\n",
      " * ResNet20 Prec@1 62.050\n",
      "\n",
      "===> epoch: 3/200\n",
      "Training:\n",
      "[0/391]\tTime 0.308 (0.308)\tData 0.020 (0.020)\tLoss 0.5746 (0.5746)\tPrec@1 (71.094)\n",
      "[40/391]\tTime 0.301 (0.292)\tData 0.020 (0.019)\tLoss 0.6375 (0.6685)\tPrec@1 (64.615)\n",
      "[80/391]\tTime 0.286 (0.305)\tData 0.019 (0.029)\tLoss 0.4912 (0.6422)\tPrec@1 (66.570)\n",
      "[120/391]\tTime 0.288 (0.301)\tData 0.020 (0.026)\tLoss 0.5843 (0.6334)\tPrec@1 (66.916)\n",
      "[160/391]\tTime 0.295 (0.300)\tData 0.019 (0.025)\tLoss 0.5981 (0.6256)\tPrec@1 (67.377)\n",
      "[200/391]\tTime 0.287 (0.300)\tData 0.018 (0.024)\tLoss 0.4950 (0.6110)\tPrec@1 (68.012)\n",
      "[240/391]\tTime 0.300 (0.299)\tData 0.023 (0.023)\tLoss 0.4883 (0.6061)\tPrec@1 (68.192)\n",
      "[280/391]\tTime 0.313 (0.299)\tData 0.019 (0.022)\tLoss 0.4834 (0.6012)\tPrec@1 (68.466)\n",
      "[320/391]\tTime 0.311 (0.299)\tData 0.025 (0.022)\tLoss 0.5215 (0.5949)\tPrec@1 (68.728)\n",
      "[360/391]\tTime 0.306 (0.299)\tData 0.019 (0.022)\tLoss 0.4398 (0.5879)\tPrec@1 (68.962)\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.019 (0.019)\tLoss 1.2049 (1.2049)\tPrec@1 54.688 (54.688)\n",
      "Test: [40/79]\tTime 0.019 (0.020)\tLoss 1.1635 (1.2018)\tPrec@1 58.594 (60.385)\n",
      " * ResNet20 Prec@1 59.860\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.0899 (1.0899)\tPrec@1 57.812 (57.812)\n",
      "Test: [40/79]\tTime 0.019 (0.020)\tLoss 1.4189 (1.1957)\tPrec@1 50.781 (61.014)\n",
      " * ResNet20 Prec@1 60.640\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.019 (0.019)\tLoss 0.9568 (0.9568)\tPrec@1 65.625 (65.625)\n",
      "Test: [40/79]\tTime 0.019 (0.020)\tLoss 0.9563 (0.9512)\tPrec@1 65.625 (68.140)\n",
      " * ResNet20 Prec@1 67.720\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 0.8087 (0.8087)\tPrec@1 74.219 (74.219)\n",
      "Test: [40/79]\tTime 0.020 (0.021)\tLoss 0.8701 (0.8518)\tPrec@1 70.312 (71.265)\n",
      " * ResNet20 Prec@1 70.300\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 0.8872 (0.8872)\tPrec@1 66.406 (66.406)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 0.9204 (0.9059)\tPrec@1 64.844 (68.521)\n",
      " * ResNet20 Prec@1 68.440\n",
      "\n",
      "===> epoch: 4/200\n",
      "Training:\n",
      "[0/391]\tTime 0.308 (0.308)\tData 0.021 (0.021)\tLoss 0.5368 (0.5368)\tPrec@1 (69.531)\n",
      "[40/391]\tTime 0.312 (0.300)\tData 0.020 (0.020)\tLoss 0.5565 (0.5236)\tPrec@1 (72.008)\n",
      "[80/391]\tTime 0.287 (0.295)\tData 0.021 (0.020)\tLoss 0.5738 (0.5191)\tPrec@1 (72.193)\n",
      "[120/391]\tTime 0.293 (0.295)\tData 0.019 (0.020)\tLoss 0.5307 (0.5199)\tPrec@1 (72.198)\n",
      "[160/391]\tTime 0.294 (0.295)\tData 0.019 (0.020)\tLoss 0.4480 (0.5161)\tPrec@1 (72.472)\n",
      "[200/391]\tTime 0.288 (0.296)\tData 0.019 (0.020)\tLoss 0.4441 (0.5115)\tPrec@1 (72.750)\n",
      "[240/391]\tTime 0.296 (0.298)\tData 0.020 (0.020)\tLoss 0.5028 (0.5060)\tPrec@1 (72.925)\n",
      "[280/391]\tTime 0.309 (0.298)\tData 0.020 (0.020)\tLoss 0.4346 (0.4985)\tPrec@1 (73.221)\n",
      "[320/391]\tTime 0.290 (0.298)\tData 0.019 (0.020)\tLoss 0.5627 (0.4995)\tPrec@1 (73.267)\n",
      "[360/391]\tTime 0.304 (0.299)\tData 0.020 (0.020)\tLoss 0.4450 (0.4984)\tPrec@1 (73.251)\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.2271 (1.2271)\tPrec@1 57.812 (57.812)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 1.2800 (1.2813)\tPrec@1 57.031 (57.393)\n",
      " * ResNet20 Prec@1 56.930\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.0057 (1.0057)\tPrec@1 60.938 (60.938)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 0.9795 (1.0113)\tPrec@1 60.156 (65.777)\n",
      " * ResNet20 Prec@1 64.430\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 0.7280 (0.7280)\tPrec@1 74.219 (74.219)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 0.8062 (0.8240)\tPrec@1 72.656 (70.389)\n",
      " * ResNet20 Prec@1 70.380\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 0.8224 (0.8224)\tPrec@1 72.656 (72.656)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 0.8687 (0.8326)\tPrec@1 67.188 (70.770)\n",
      " * ResNet20 Prec@1 70.530\n",
      "Testing:\n",
      "Test: [0/79]\tTime 0.020 (0.020)\tLoss 1.0995 (1.0995)\tPrec@1 60.938 (60.938)\n",
      "Test: [40/79]\tTime 0.020 (0.020)\tLoss 1.0901 (1.1271)\tPrec@1 59.375 (61.528)\n",
      " * ResNet20 Prec@1 61.670\n",
      "\n",
      "===> epoch: 5/200\n",
      "Training:\n",
      "[0/391]\tTime 0.309 (0.309)\tData 0.021 (0.021)\tLoss 0.5283 (0.5283)\tPrec@1 (71.875)\n",
      "[40/391]\tTime 0.306 (0.300)\tData 0.019 (0.020)\tLoss 0.4442 (0.4384)\tPrec@1 (75.038)\n",
      "[80/391]\tTime 0.289 (0.300)\tData 0.020 (0.020)\tLoss 0.4321 (0.4324)\tPrec@1 (75.463)\n",
      "[120/391]\tTime 0.308 (0.301)\tData 0.020 (0.020)\tLoss 0.5071 (0.4345)\tPrec@1 (75.549)\n",
      "[160/391]\tTime 0.293 (0.301)\tData 0.020 (0.020)\tLoss 0.3852 (0.4380)\tPrec@1 (75.582)\n",
      "[200/391]\tTime 0.292 (0.302)\tData 0.020 (0.020)\tLoss 0.3895 (0.4433)\tPrec@1 (75.272)\n",
      "[240/391]\tTime 0.291 (0.301)\tData 0.020 (0.020)\tLoss 0.2974 (0.4396)\tPrec@1 (75.486)\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "import os\n",
    "import time\n",
    "import logging\n",
    "import argparse\n",
    "import numpy as np\n",
    "from visdom import Visdom\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "from utils import *\n",
    "\n",
    "# Teacher models:\n",
    "# VGG11/VGG13/VGG16/VGG19, GoogLeNet, AlxNet, ResNet18, ResNet34, \n",
    "# ResNet50, ResNet101, ResNet152, ResNeXt29_2x64d, ResNeXt29_4x64d, \n",
    "# ResNeXt29_8x64d, ResNeXt29_32x64d, PreActResNet18, PreActResNet34, \n",
    "# PreActResNet50, PreActResNet101, PreActResNet152, \n",
    "# DenseNet121, DenseNet161, DenseNet169, DenseNet201, \n",
    "import models\n",
    "\n",
    "# Student models:\n",
    "# myNet, LeNet, FitNet\n",
    "\n",
    "start_time = time.time()\n",
    "# os.makedirs('./checkpoint', exist_ok=True)\n",
    "\n",
    "# Training settings\n",
    "parser = argparse.ArgumentParser(description='PyTorch DML 3S')\n",
    "\n",
    "parser.add_argument('--dataset',\n",
    "                    choices=['CIFAR10',\n",
    "                             'CIFAR100'\n",
    "                            ],\n",
    "                    default='CIFAR10')\n",
    "parser.add_argument('--net1',\n",
    "                    choices=['ResNet20',\n",
    "                             'ResNet32',\n",
    "                             'ResNet50',\n",
    "                             'ResNet56',\n",
    "                             'ResNet110'\n",
    "                            ],\n",
    "                    default='ResNet20')\n",
    "parser.add_argument('--net2',\n",
    "                    choices=['ResNet20',\n",
    "                             'ResNet32',\n",
    "                             'ResNet50',\n",
    "                             'ResNet56',\n",
    "                             'ResNet110'\n",
    "                            ],\n",
    "                    default='ResNet20')\n",
    "\n",
    "parser.add_argument('--net3',\n",
    "                    choices=['ResNet20',\n",
    "                             'ResNet32',\n",
    "                             'ResNet50',\n",
    "                             'ResNet56',\n",
    "                             'ResNet110'\n",
    "                            ],\n",
    "                    default='ResNet20')\n",
    "\n",
    "parser.add_argument('--net4',\n",
    "                    choices=['ResNet20',\n",
    "                             'ResNet32',\n",
    "                             'ResNet34',\n",
    "                             'ResNet50',\n",
    "                             'ResNet56',\n",
    "                             'ResNet110',\n",
    "                             'VGG19',\n",
    "                             'GoogleNet',\n",
    "                             'DenseNet121'\n",
    "                            ],\n",
    "                    default='ResNet20')\n",
    "\n",
    "parser.add_argument('--net5',\n",
    "                    choices=['ResNet20',\n",
    "                             'ResNet32',\n",
    "                             'ResNet34',\n",
    "                             'ResNet50',\n",
    "                             'ResNet56',\n",
    "                             'ResNet110',\n",
    "                             'VGG19',\n",
    "                             'GoogleNet',\n",
    "                             'DenseNet121'\n",
    "                            ],\n",
    "                    default='ResNet20')\n",
    "\n",
    "parser.add_argument('--n_class', type=int, default=10, metavar='N', help='num of classes')\n",
    "parser.add_argument('--batch_size', type=int, default=128, metavar='N', help='input batch size for training')\n",
    "parser.add_argument('--test_batch_size', type=int, default=128, metavar='N', help='input test batch size for training')\n",
    "parser.add_argument('--epochs', type=int, default=20, metavar='N', help='number of epochs to train (default: 20)')\n",
    "parser.add_argument('--lr', type=float, default=0.1, metavar='LR', help='learning rate (default: 0.01)')\n",
    "parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='SGD momentum (default: 0.5)')\n",
    "parser.add_argument('--device', default='cuda:1', type=str, help='device: cuda or cpu')\n",
    "parser.add_argument('--print_freq', type=int, default=40, metavar='N', help='how many batches to wait before logging training status')\n",
    "\n",
    "config = ['--dataset', 'CIFAR10', '--epochs', '200', '--n_class', '10', '--device', 'cuda:0']\n",
    "args = parser.parse_args(config)\n",
    "\n",
    "device = args.device if torch.cuda.is_available() else 'cpu'\n",
    "save_dir = './checkpoint/' + args.dataset + '/'\n",
    "\n",
    "# models\n",
    "net1 = getattr(models, args.net1)(num_classes=args.n_class)\n",
    "net1.to(device)\n",
    "net2 = getattr(models, args.net2)(num_classes=args.n_class)\n",
    "net2.to(device)\n",
    "net3 = getattr(models, args.net3)(num_classes=args.n_class)\n",
    "net3.to(device)\n",
    "net4 = getattr(models, args.net4)(num_classes=args.n_class)\n",
    "net4.to(device)\n",
    "net5 = getattr(models, args.net5)(num_classes=args.n_class)\n",
    "net5.to(device)\n",
    "\n",
    "# logging\n",
    "logfile = save_dir + 'DML_3S_.log'\n",
    "if os.path.exists(logfile):\n",
    "    os.remove(logfile)\n",
    "def log_out(info):\n",
    "    f = open(logfile, mode='a')\n",
    "    f.write(info)\n",
    "    f.write('\\n')\n",
    "    f.close()\n",
    "    print(info)\n",
    "    \n",
    "# visualizer\n",
    "vis = Visdom(env='distill')\n",
    "loss_win = vis.line(\n",
    "    X=np.column_stack((0,0,0,0,0)),\n",
    "    Y=np.column_stack((0,0,0,0,0)),\n",
    "    opts=dict(\n",
    "        title='DML 5S Loss',\n",
    "        xlabel='epoch',\n",
    "        xtickmin=0,\n",
    "        ylabel='loss',\n",
    "        ytickmin=0,\n",
    "        ytickstep=0.5,\n",
    "        legend=['net1_loss', 'net2_loss', 'net3_loss', 'net4_loss', 'net5_loss']\n",
    "    ),\n",
    "    name=\"loss\"\n",
    ")\n",
    "\n",
    "acc_win = vis.line(\n",
    "    X=np.column_stack((0,0,0,0,0)),\n",
    "    Y=np.column_stack((0,0,0,0,0)),\n",
    "    opts=dict(\n",
    "        title='DML 5S Acc',\n",
    "        xlabel='epoch',\n",
    "        xtickmin=0,\n",
    "        ylabel='accuracy',\n",
    "        ytickmin=0,\n",
    "        ytickmax=100,\n",
    "        legend=['net1_acc', 'net2_acc', 'net3_acc', 'net4_acc', 'net5_acc']\n",
    "    ),\n",
    "    name=\"acc\"\n",
    ")\n",
    "\n",
    "\n",
    "# data\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomCrop(32, 4),\n",
    "    transforms.ToTensor(),\n",
    "    normalize,\n",
    "])\n",
    "test_transform = transforms.Compose([transforms.ToTensor(), normalize])\n",
    "train_set = getattr(datasets, args.dataset)(root='../data', train=True, download=True, transform=train_transform)\n",
    "test_set = getattr(datasets, args.dataset)(root='../data', train=False, download=False, transform=test_transform)\n",
    "train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_set, batch_size=args.test_batch_size, shuffle=False)\n",
    "\n",
    "# optimizer = optim.SGD(st_model.parameters(), lr=args.lr, momentum=args.momentum)\n",
    "optimizer_1 = optim.SGD(net1.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)\n",
    "optimizer_2 = optim.SGD(net2.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)\n",
    "optimizer_3 = optim.SGD(net3.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)\n",
    "optimizer_4 = optim.SGD(net4.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)\n",
    "optimizer_5 = optim.SGD(net5.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)\n",
    "\n",
    "lr_scheduler_1 = optim.lr_scheduler.MultiStepLR(optimizer_1, milestones=[100, 150])\n",
    "lr_scheduler_2 = optim.lr_scheduler.MultiStepLR(optimizer_2, milestones=[100, 150])\n",
    "lr_scheduler_3 = optim.lr_scheduler.MultiStepLR(optimizer_3, milestones=[100, 150])\n",
    "lr_scheduler_4 = optim.lr_scheduler.MultiStepLR(optimizer_4, milestones=[100, 150])\n",
    "lr_scheduler_5 = optim.lr_scheduler.MultiStepLR(optimizer_5, milestones=[100, 150])\n",
    "\n",
    "\n",
    "# train with multi-teacher\n",
    "def train(epoch, net1, net2, net3, net4, net5):\n",
    "    print('Training:')\n",
    "    # switch to train mode\n",
    "    net1.train()\n",
    "    net2.train()\n",
    "    net3.train()\n",
    "    net4.train()\n",
    "    net5.train()\n",
    "    batch_time = AverageMeter()\n",
    "    data_time = AverageMeter()\n",
    "    losses_1 = AverageMeter()\n",
    "    losses_2 = AverageMeter()\n",
    "    losses_3 = AverageMeter()\n",
    "    losses_4 = AverageMeter()\n",
    "    losses_5 = AverageMeter()\n",
    "    top1_1 = AverageMeter()\n",
    "    top1_2 = AverageMeter()\n",
    "    top1_3 = AverageMeter()\n",
    "    top1_4 = AverageMeter()\n",
    "    top1_5 = AverageMeter()\n",
    "    \n",
    "    end = time.time()\n",
    "    for i, (input, target) in enumerate(train_loader):\n",
    "\n",
    "        # measure data loading time\n",
    "        data_time.update(time.time() - end)\n",
    "\n",
    "        input, target = input.to(device), target.to(device)\n",
    "        \n",
    "        # compute outputs\n",
    "        _,_,_,_, output_1 = net1(input)\n",
    "        _,_,_,_, output_2 = net2(input)\n",
    "        _,_,_,_, output_3 = net3(input)\n",
    "        _,_,_,_, output_4 = net4(input)\n",
    "        _,_,_,_, output_5 = net5(input)\n",
    "        logits_1 = F.softmax(output_1)\n",
    "        logits_2 = F.softmax(output_2)\n",
    "        logits_3 = F.softmax(output_3)\n",
    "        logits_4 = F.softmax(output_4)\n",
    "        logits_5 = F.softmax(output_5)\n",
    "        \n",
    "        optimizer_1.zero_grad()\n",
    "        loss_1 = 0.5 * (nn.KLDivLoss()(logits_2, logits_1) + nn.KLDivLoss()(logits_3, logits_1) \n",
    "                 + nn.KLDivLoss()(logits_4, logits_1)+ nn.KLDivLoss()(logits_5, logits_1)) + F.cross_entropy(output_1, target)\n",
    "        loss_1.backward(retain_graph=True)\n",
    "        optimizer_1.step()\n",
    "        \n",
    "        optimizer_2.zero_grad()\n",
    "        loss_2 = 0.5 * (nn.KLDivLoss()(logits_1, logits_2) + nn.KLDivLoss()(logits_3, logits_2)\n",
    "                 + nn.KLDivLoss()(logits_4, logits_2)+ nn.KLDivLoss()(logits_5, logits_2))+ F.cross_entropy(output_2, target)\n",
    "        loss_2.backward(retain_graph=True)\n",
    "        optimizer_2.step()\n",
    "        \n",
    "        optimizer_3.zero_grad()\n",
    "        loss_3 = 0.5 * (nn.KLDivLoss()(logits_1, logits_3) + nn.KLDivLoss()(logits_2, logits_3) \n",
    "                 + nn.KLDivLoss()(logits_4, logits_3)+ nn.KLDivLoss()(logits_5, logits_3))+ F.cross_entropy(output_3, target)\n",
    "        loss_3.backward(retain_graph=True)\n",
    "        optimizer_3.step()\n",
    "        \n",
    "        optimizer_4.zero_grad()\n",
    "        loss_4 = 0.5 * (nn.KLDivLoss()(logits_1, logits_4) + nn.KLDivLoss()(logits_2, logits_4)\n",
    "                 + nn.KLDivLoss()(logits_3, logits_4)+ nn.KLDivLoss()(logits_5, logits_4))+ F.cross_entropy(output_4, target)\n",
    "        loss_4.backward(retain_graph=True)\n",
    "        optimizer_4.step()\n",
    "        \n",
    "        optimizer_5.zero_grad()\n",
    "        loss_5 = 0.5 * (nn.KLDivLoss()(logits_1, logits_5) + nn.KLDivLoss()(logits_2, logits_5)\n",
    "                 + nn.KLDivLoss()(logits_3, logits_5)+ nn.KLDivLoss()(logits_4, logits_5))+ F.cross_entropy(output_5, target)\n",
    "        loss_5.backward(retain_graph=True)\n",
    "        optimizer_5.step()\n",
    "\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        net1_acc = accuracy(output_1, target)[0]\n",
    "        net2_acc = accuracy(output_2, target)[0]\n",
    "        net3_acc = accuracy(output_3, target)[0]\n",
    "        net4_acc = accuracy(output_4, target)[0]\n",
    "        net5_acc = accuracy(output_5, target)[0]\n",
    "        losses_1.update(loss_1.item(), input.size(0))\n",
    "        top1_1.update(net1_acc, input.size(0))\n",
    "        losses_2.update(loss_2.item(), input.size(0))\n",
    "        top1_2.update(net2_acc, input.size(0))\n",
    "        losses_3.update(loss_3.item(), input.size(0))\n",
    "        top1_3.update(net3_acc, input.size(0))\n",
    "        losses_4.update(loss_4.item(), input.size(0))\n",
    "        top1_4.update(net4_acc, input.size(0))\n",
    "        losses_5.update(loss_5.item(), input.size(0))\n",
    "        top1_5.update(net5_acc, input.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        if i % args.print_freq == 0:\n",
    "            log_out('[{0}/{1}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                  'Prec@1 ({top1_1.avg:.3f})'.format(\n",
    "                      i, len(train_loader), batch_time=batch_time,\n",
    "                      data_time=data_time, loss=losses_1, top1_1=top1_1))\n",
    "    return losses_1.avg, losses_2.avg, losses_3.avg, losses_4.avg, losses_5.avg, net1_acc.cpu().numpy(), net2_acc.cpu().numpy(), net3_acc.cpu().numpy(), net4_acc.cpu().numpy(), net5_acc.cpu().numpy()\n",
    "\n",
    "\n",
    "def test(model):\n",
    "    print('Testing:')\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "    batch_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "\n",
    "    end = time.time()\n",
    "    with torch.no_grad():\n",
    "        for i, (input, target) in enumerate(test_loader):\n",
    "            input, target = input.to(device), target.to(device)\n",
    "\n",
    "            # compute output\n",
    "            _,_,_,_,output = model(input)\n",
    "            loss = F.cross_entropy(output, target)\n",
    "\n",
    "            output = output.float()\n",
    "            loss = loss.float()\n",
    "\n",
    "            # measure accuracy and record loss\n",
    "            test_acc = accuracy(output.data, target.data)[0]\n",
    "            losses.update(loss.item(), input.size(0))\n",
    "            top1.update(test_acc, input.size(0))\n",
    "\n",
    "            # measure elapsed time\n",
    "            batch_time.update(time.time() - end)\n",
    "            end = time.time()\n",
    "\n",
    "            if i % args.print_freq == 0:\n",
    "                log_out('Test: [{0}/{1}]\\t'\n",
    "                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                      'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'\n",
    "                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                          i, len(test_loader), batch_time=batch_time, loss=losses,\n",
    "                          top1=top1))\n",
    "\n",
    "    log_out(' * {0} Prec@1 {top1.avg:.3f}'.format(model.model_name, top1=top1))\n",
    "\n",
    "    return losses.avg, test_acc.cpu().numpy(), top1.avg.cpu().numpy()\n",
    "\n",
    "\n",
    "print('*-----------------DML----------------*')\n",
    "best_acc = 0\n",
    "for epoch in range(1, args.epochs + 1):\n",
    "    log_out(\"\\n===> epoch: {}/{}\".format(epoch, args.epochs))\n",
    "#     log_out('current lr {:.5e}'.format(optimizer_1.param_groups[0]['lr']))\n",
    "    lr_scheduler_1.step()\n",
    "    lr_scheduler_2.step()\n",
    "    lr_scheduler_3.step()\n",
    "    lr_scheduler_4.step()\n",
    "    lr_scheduler_5.step()\n",
    "    train_loss_1, train_loss_2, train_loss_3, train_loss_4, train_loss_5, net1_acc, net2_acc, net3_acc, net4_acc, net5_acc = train(epoch, net1, net2, net3, net4, net5)\n",
    "    # visaulize loss\n",
    "    vis.line(np.column_stack((train_loss_1, train_loss_2, train_loss_3, train_loss_4, train_loss_5)), np.column_stack((epoch, epoch, epoch, epoch, epoch)), loss_win, update=\"append\")\n",
    "    _, test_acc_1, top1_1 = test(net1)\n",
    "    _, test_acc_2, top1_2 = test(net2)\n",
    "    _, test_acc_3, top1_3 = test(net3)\n",
    "    _, test_acc_4, top1_4 = test(net4)\n",
    "    _, test_acc_5, top1_5 = test(net5)\n",
    "    vis.line(np.column_stack((top1_1, top1_2, top1_3, top1_4, top1_5)), np.column_stack((epoch, epoch, epoch, epoch, epoch)), acc_win, update=\"append\")\n",
    "    best_acc = max(top1_1, best_acc)\n",
    "\n",
    "log_out(\"@ BEST Prec: {:.4f}\".format(best_acc))\n",
    "log_out(\"--- {:.3f} mins ---\".format((time.time() - start_time)/60))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py37] *",
   "language": "python",
   "name": "conda-env-py37-py"
  },
  "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": 2
}