adaptive_FitNet_teacher-level.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-19T14:16:34.572Z"
    }
   },
   "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",
      "StudentNet:\n",
      "\n",
      "ResNet(\n",
      "  (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "  (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (layer1): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (2): BasicBlock(\n",
      "      (conv1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (shortcut): LambdaLayer()\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (2): BasicBlock(\n",
      "      (conv1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\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",
      "      (shortcut): LambdaLayer()\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",
      "      (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",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (2): 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",
      "      (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",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "  )\n",
      "  (linear): Linear(in_features=64, out_features=10, bias=True)\n",
      ")\n",
      "Training adapter:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:246: 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:209: 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/_reduction.py:15: UserWarning: reduction='elementwise_mean' is deprecated, please use reduction='mean' instead.\n",
      "  warnings.warn(\"reduction='elementwise_mean' is deprecated, please use reduction='mean' instead.\")\n",
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:180: UserWarning: Implicit dimension choice for log_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"
     ]
    }
   ],
   "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",
    "from PIL import Image\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",
    "from metric.loss import FitNet, AttentionTransfer, RKdAngle, RkdDistance\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",
    "# PreActhttps://www.bing.com/?mkt=zh-CNResNet50, 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 ada. FitNet')\n",
    "\n",
    "parser.add_argument('--dataset',\n",
    "                    choices=['CIFAR10',\n",
    "                             'CIFAR100'\n",
    "                            ],\n",
    "                    default='CIFAR10')\n",
    "parser.add_argument('--teachers',\n",
    "                    choices=['ResNet32',\n",
    "                             'ResNet50',\n",
    "                             'ResNet56',\n",
    "                             'ResNet110',\n",
    "                             'DenseNet121'\n",
    "                            ],\n",
    "                    default=['ResNet32', 'ResNet56', 'ResNet110'],\n",
    "                    nargs='+')\n",
    "parser.add_argument('--student',\n",
    "                    choices=['ResNet20',\n",
    "                             'myNet'\n",
    "                            ],\n",
    "                    default='ResNet20')\n",
    "\n",
    "parser.add_argument('--kd_ratio', default=0.7, type=float)\n",
    "parser.add_argument('--n_class', type=int, default=10, metavar='N', help='num of classes')\n",
    "parser.add_argument('--T', type=float, default=20.0, metavar='Temputure', help='Temputure for distillation')\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:0', type=str, help='device: cuda or cpu')\n",
    "parser.add_argument('--print_freq', type=int, default=10, metavar='N', help='how many batches to wait before logging training status')\n",
    "\n",
    "config = ['--epochs', '200', '--T', '5.0', '--device', 'cuda:0']\n",
    "args = parser.parse_args(config)\n",
    "\n",
    "device = args.device if torch.cuda.is_available() else 'cpu'\n",
    "load_dir = './checkpoint/' + args.dataset + '/'\n",
    "\n",
    "# teachers model\n",
    "teacher_models = []\n",
    "for te in args.teachers:\n",
    "    te_model = getattr(models, te)(num_classes=args.n_class)\n",
    "#     print(te_model)\n",
    "    te_model.load_state_dict(torch.load(load_dir + te_model.model_name + '.pth'))\n",
    "    te_model.to(device)\n",
    "    teacher_models.append(te_model)\n",
    "\n",
    "st_model = getattr(models, args.student)(num_classes=args.n_class)  # args.student()\n",
    "st_model.to(device)\n",
    "\n",
    "# logging\n",
    "logfile = load_dir + 'ada_te_fitnet_' + st_model.model_name + '.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.array([0]),\n",
    "    Y=np.array([0]),\n",
    "    opts=dict(\n",
    "        title='FitNet ada. loss',\n",
    "        xtickmin=0,\n",
    "#         xtickmax=1,\n",
    "#         xtickstep=5,\n",
    "        ytickmin=0,\n",
    "#         ytickmax=1,\n",
    "        ytickstep=0.5,\n",
    "#         markers=True,\n",
    "#         markersymbol='dot',\n",
    "#         markersize=5,\n",
    "    ),\n",
    "    name=\"loss\"\n",
    ")\n",
    "\n",
    "acc_win = vis.line(\n",
    "    X=np.column_stack((0, 0)),\n",
    "    Y=np.column_stack((0, 0)),\n",
    "    opts=dict(\n",
    "        title='FitNet ada. ACC',\n",
    "        xtickmin=0,\n",
    "#         xtickstep=5,\n",
    "        ytickmin=0,\n",
    "        ytickmax=100,\n",
    "#         markers=True,\n",
    "#         markersymbol='dot',\n",
    "#         markersize=5,\n",
    "        legend=['train_acc', 'test_acc']\n",
    "    ),\n",
    "    name=\"acc\"\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",
    "# optim\n",
    "optimizer_W = optim.SGD([adapter.W], lr=args.lr, momentum=0.9)\n",
    "optimizer_theta = optim.SGD([adapter.theta], lr=args.lr, momentum=0.9)\n",
    "optimizer_sgd = optim.SGD(st_model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)\n",
    "lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer_sgd, gamma=0.1, milestones=[100, 150])\n",
    "lr_scheduler2 = optim.lr_scheduler.MultiStepLR(optimizer_W, milestones=[40, 50])\n",
    "lr_scheduler3 = optim.lr_scheduler.MultiStepLR(optimizer_theta, milestones=[40, 50])\n",
    "\n",
    "# loss\n",
    "def kd_criterion(y, labels, weighted_logits, T=10.0, alpha=0.7):\n",
    "    return nn.KLDivLoss()(F.log_softmax(y/T), weighted_logits) * (T*T * 2.0 * alpha) + F.cross_entropy(y, labels) * (1. - alpha)\n",
    "\n",
    "dist_criterion = RkdDistance().to(device)\n",
    "angle_criterion = RKdAngle().to(device)\n",
    "fitnet_criterion = [FitNet(32, 64), FitNet(64, 64),FitNet(64, 64)]\n",
    "[f.to(device) for f in fitnet_criterion]\n",
    "\n",
    "# adapter model\n",
    "class Adapter():\n",
    "    def __init__(self, in_models, pool_size):\n",
    "        # representations of teachers\n",
    "        pool_ch = pool_size[1]  # 64\n",
    "        pool_w = pool_size[2]   # 8\n",
    "        LR_list = []\n",
    "        torch.manual_seed(1)\n",
    "        self.theta = torch.randn(len(in_models), pool_ch).to(device)  # [3, 64]\n",
    "        self.theta.requires_grad_(True)\n",
    "   \n",
    "        self.max_feat = nn.MaxPool2d(kernel_size=(pool_w, pool_w), stride=pool_w).to(device)\n",
    "        self.W = torch.randn(pool_ch, 1).to(device)\n",
    "        self.W.requires_grad_(True)\n",
    "        self.val = False\n",
    "\n",
    "    def loss(self, y, labels, weighted_logits, T=10.0, alpha=0.7):\n",
    "        ls = nn.KLDivLoss()(F.log_softmax(y/T), weighted_logits) * (T*T * 2.0 * alpha) + F.cross_entropy(y, labels) * (1. - alpha)\n",
    "        if not self.val:\n",
    "            ls += 0.1 * (torch.sum(self.W * self.W) + torch.sum(torch.sum(self.theta * self.theta, dim=1), dim=0))\n",
    "        return ls\n",
    "        \n",
    "    def gradient(self, lr=0.01):\n",
    "        self.W.data = self.W.data - lr * self.W.grad.data\n",
    "        # Manually zero the gradients after updating weights\n",
    "        self.W.grad.data.zero_()\n",
    "        \n",
    "    def eval(self):\n",
    "        self.val = True\n",
    "        self.theta.detach()\n",
    "        self.W.detach()\n",
    "    \n",
    "    # input size: [64, 8, 8], [128, 3, 10]\n",
    "    def forward(self, conv_map, te_logits_list):\n",
    "        beta = self.max_feat(conv_map)\n",
    "        beta = torch.squeeze(beta)  # [128, 64]\n",
    "        \n",
    "        latent_factor = []\n",
    "        for t in self.theta:\n",
    "            latent_factor.append(beta * t)\n",
    "#         latent_factor = torch.stack(latent_factor, dim=0)  # [3, 128, 64]\n",
    "        alpha = []\n",
    "        for lf in latent_factor:  # lf.size:[128, 64]\n",
    "            alpha.append(lf.mm(self.W))\n",
    "        alpha = torch.stack(alpha, dim=0)  # [3, 128, 1]\n",
    "        alpha = torch.squeeze(alpha).transpose(0, 1) # [128, 3]\n",
    "        weight = F.softmax(alpha)  # [128, 3]\n",
    "\n",
    "        return weight\n",
    "\n",
    "# adapter instance\n",
    "_,_,_,pool_m,_ = st_model(torch.randn(1,3, 128, 128).to(device))  # get pool_size of student\n",
    "# reate adapter instance\n",
    "adapter = Adapter(teacher_models, pool_m.size())\n",
    "\n",
    "\n",
    "def teacher_weights(n_epochs=50, model=st_model):\n",
    "    print('Training adapter:')\n",
    "    start_time = time.time()\n",
    "    model.train()\n",
    "    adapter.eval()\n",
    "    for ep in range(n_epochs):\n",
    "        lr_scheduler2.step()\n",
    "        lr_scheduler3.step()\n",
    "        for i, (input, target) in enumerate(train_loader):\n",
    "\n",
    "            input, target = input.to(device), target.to(device)\n",
    "            # compute outputs\n",
    "            b1, b2, b3, pool, output = model(input) # out_feat: 16, 32, 64, 64, - \n",
    "            st_maps = [b1, b2, b3, pool]\n",
    "#             print('b1:{}, b2:{}, b3{}, pool:{}'.format(b1.size(), b2.size(), b3.size(), pool.size()))\n",
    "# b1:torch.Size([128, 16, 32, 32]), b2:torch.Size([128, 32, 16, 16]), b3torch.Size([128, 64, 8, 8]), pool:torch.Size([128, 64, 1, 1])\n",
    "\n",
    "            te_scores_list = []\n",
    "            hint_maps = []\n",
    "            fit_loss = 0\n",
    "            for j,te in enumerate(teacher_models):\n",
    "                te.eval()\n",
    "                with torch.no_grad():\n",
    "                    t_b1, t_b2, t_b3, t_pool, t_output = te(input)\n",
    "#                 print('t_b1:{}, t_b2:{}, t_b3:{}, t_pool:{}'.format(t_b1.size(), t_b2.size(), t_b3.size(), t_pool.size()))\n",
    "# t_b1:torch.Size([128, 16, 32, 32]), t_b2:torch.Size([128, 32, 16, 16]), t_b3:torch.Size([128, 64, 8, 8]), t_pool:torch.Size([128, 64, 1, 1])\n",
    "                hint_maps.append(t_pool)\n",
    "                t_output = F.softmax(t_output/args.T)\n",
    "                te_scores_list.append(t_output)\n",
    "            te_scores_Tensor = torch.stack(te_scores_list, dim=1)  # size: [128, 3, 10]\n",
    "            \n",
    "            weight = adapter.forward(pool, te_scores_Tensor)\n",
    "            weight_t = torch.unsqueeze(weight, dim=2)  # [128, 3, 1]\n",
    "            weighted_logits = weight_t * te_scores_Tensor  # [128, 3, 10]\n",
    "            weighted_logits = torch.sum(weighted_logits, dim=1)\n",
    "            \n",
    "            optimizer_sgd.zero_grad()\n",
    "            optimizer_W.zero_grad()\n",
    "            optimizer_theta.zero_grad()\n",
    "            \n",
    "            angle_loss = angle_criterion(output, weighted_logits)\n",
    "            dist_loss = dist_criterion(output, weighted_logits)\n",
    "            # compute gradient and do SGD step\n",
    "            ada_loss = adapter.loss(output, target, weighted_logits, T=args.T, alpha=args.kd_ratio)\n",
    "            \n",
    "            for j in range(len(teacher_models)):\n",
    "                fit_loss += fitnet_criterion[j](st_maps[j+1], hint_maps[j])\n",
    "#             fit_loss = fitnet_criterion[0](b2, hint_maps[0][3]) + fitnet_criterion[1](b3, hint_maps[1][3]) + fitnet_criterion(pool, hint_maps[2][3])\n",
    "            loss = ada_loss + dist_loss + angle_loss + fit_loss\n",
    "            \n",
    "            loss.backward(retain_graph=True)\n",
    "            optimizer_sgd.step()\n",
    "            optimizer_W.step()\n",
    "            optimizer_theta.step()\n",
    "#          vis.line(np.array([loss.item()]), np.array([ep]), loss_win, update=\"append\")\n",
    "        log_out('epoch[{}/{}]adapter Loss: {:.4f}'.format(ep, n_epochs, loss.item()))\n",
    "    end_time = time.time()\n",
    "    log_out(\"--- adapter training cost {:.3f} mins ---\".format((end_time - start_time)/60))\n",
    "    return torch.mean(weight_t, dim=0)\n",
    "\n",
    "# train with multi-teacher\n",
    "def train(epoch, model, te_weights):\n",
    "    print('Training:')\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "\n",
    "    batch_time = AverageMeter()\n",
    "    data_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = 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",
    "        b1, b2, b3, pool, output = model(input)\n",
    "        st_maps = [b1, b2, b3, pool]\n",
    "        \n",
    "        te_scores_list = []\n",
    "        hint_maps = []\n",
    "        fit_loss = 0\n",
    "        for j,te in enumerate(teacher_models):\n",
    "            te.eval()\n",
    "            with torch.no_grad():\n",
    "                t_b1, t_b2, t_b3, t_pool, t_output = te(input)\n",
    "        \n",
    "            hint_maps.append(t_pool)\n",
    "            t_output = F.softmax(t_output/args.T)\n",
    "            \n",
    "            te_scores_list.append(t_output)\n",
    "        te_scores_Tensor = torch.stack(te_scores_list, dim=1)  # size: [128, 3, 10]\n",
    "        weighted_logits = te_scores_Tensor * te_weights\n",
    "        weighted_logits = torch.sum(weighted_logits, dim=1)\n",
    "        \n",
    "        optimizer_sgd.zero_grad()\n",
    "        \n",
    "        angle_loss = angle_criterion(output, weighted_logits)\n",
    "        dist_loss = dist_criterion(output, weighted_logits)\n",
    "        \n",
    "        # compute gradient and do SGD step\n",
    "        kd_loss = kd_criterion(output, target, weighted_logits, T=args.T, alpha=args.kd_ratio)\n",
    "        for j in range(len(teacher_models)):\n",
    "            fit_loss += fitnet_criterion[j](st_maps[j+1], hint_maps[j])\n",
    "        \n",
    "        loss = kd_loss # + dist_loss + angle_loss + fit_loss\n",
    "\n",
    "        loss.backward(retain_graph=True)\n",
    "        optimizer_sgd.step()\n",
    "\n",
    "        output = output.float()\n",
    "        loss = loss.float()\n",
    "        # measure accuracy and record loss\n",
    "        train_acc = accuracy(output.data, target.data)[0]\n",
    "        losses.update(loss.item(), input.size(0))\n",
    "        top1.update(train_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.val:.3f} ({top1.avg:.3f})'.format(\n",
    "                      i, len(train_loader), batch_time=batch_time,\n",
    "                      data_time=data_time, loss=losses, top1=top1))\n",
    "    return losses.avg, train_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(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))\n",
    "\n",
    "    return losses.avg, test_acc.cpu().numpy(), top1.avg.cpu().numpy()\n",
    "\n",
    "# \"\"\"\n",
    "print('StudentNet:\\n')\n",
    "print(st_model)\n",
    "st_model.apply(weights_init_normal)\n",
    "weights = teacher_weights(n_epochs=50)\n",
    "weights = weights.detach()\n",
    "log_out('------------weight:{}'.format(weights))\n",
    "# st_model.apply(weights_init_normal)\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_sgd.param_groups[0]['lr']))\n",
    "    lr_scheduler.step(epoch)\n",
    "    train_loss, train_acc = train(epoch, st_model, weights)\n",
    "    # visaulize loss\n",
    "    vis.line(np.array([train_loss]), np.array([epoch]), loss_win, update=\"append\")\n",
    "    _, test_acc, top1 = test(st_model)\n",
    "    vis.line(np.column_stack((train_acc, top1)), np.column_stack((epoch, epoch)), acc_win, update=\"append\")\n",
    "    if top1 > best_acc:\n",
    "        best_acc = top1\n",
    "            \n",
    "# release GPU memory\n",
    "torch.cuda.empty_cache()\n",
    "log_out(\"BEST ACC: {:.3f}\".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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}