multi_teacher_avg_distill.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-08-18T10:59:47.828Z"
    }
   },
   "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=100, bias=True)\n",
      ")\n",
      "\n",
      "===> epoch: 1/200\n",
      "current lr 1.00000e-01\n",
      "Training:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yaliu/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:260: 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:146: 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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/391]\tTime 0.889 (0.889)\tData 0.041 (0.041)\tLoss 2.2553 (2.2553)\tPrec@1 0.781 (0.781)\n"
     ]
    }
   ],
   "source": [
    "### from __future__ import print_function\n",
    "import os\n",
    "import time\n",
    "import logging\n",
    "import argparse\n",
    "import random\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 multi_teacher_avg_distill')\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",
    "                            ],\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('--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:1', 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 = ['--dataset', 'CIFAR100', '--n_class', '100', '--epochs', '200', '--T', '5.0', '--device', 'cuda:1']\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",
    "    te_model.eval()  # eval mode\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 + 'avg_distill_' + 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='multi avg. Loss',\n",
    "        xlabel='epoch',\n",
    "        xtickmin=0,\n",
    "        ylabel='loss',\n",
    "        ytickmin=0,\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='multi-KD avg. Acc',\n",
    "        xlabel='epoch',\n",
    "        xtickmin=0,\n",
    "        ylabel='accuracy',\n",
    "        ytickmin=0,\n",
    "        ytickmax=100,\n",
    "        legend=['train_acc', 'test_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 = optim.Adam(st_model.parameters(), lr=args.lr)\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, milestones=[100, 150])\n",
    "\n",
    "# avg diatill\n",
    "def distillation_loss(y, labels, logits, T, alpha=0.7):\n",
    "    return nn.KLDivLoss()(F.log_softmax(y/T), logits) * (T*T * 2.0 * alpha) + F.cross_entropy(y, labels) * (1. - alpha)\n",
    "\n",
    "# triplet loss\n",
    "triplet_loss = nn.TripletMarginLoss(margin=0.2, p=2).to(device)\n",
    "\n",
    "# get max infoentropy scores\n",
    "# input: Tensor[3, 128, 10]\n",
    "def maxInfo_logits(te_scores_Tensor):\n",
    "    used_score = torch.FloatTensor(te_scores_Tensor.size(1), te_scores_Tensor.size(2)).to(device)\n",
    "    ents = torch.FloatTensor(te_scores_Tensor.size(0), te_scores_Tensor.size(1)).to(device)\n",
    "    logp = torch.log2(te_scores_Tensor)\n",
    "    plogp = -logp.mul(te_scores_Tensor)\n",
    "    for i,te in enumerate(plogp):\n",
    "        ents[i] = torch.sum(te, dim=1)\n",
    "    max_ent_index = torch.max(ents, dim=0).indices   # 取每一列最大值index\n",
    "#     print(max_ent_index)\n",
    "    for i in range(max_ent_index.size(0)):\n",
    "        used_score[i] = te_scores_Tensor[max_ent_index[i].item()][i]\n",
    "#     print(used_score)\n",
    "\n",
    "    return used_score\n",
    "    \n",
    "# avg logits\n",
    "# input: Tensor[3, 128, 10]\n",
    "def avg_logits(te_scores_Tensor):\n",
    "#     print(te_scores_Tensor.size())\n",
    "    mean_Tensor = torch.mean(te_scores_Tensor, dim=1)\n",
    "#     print(mean_Tensor)\n",
    "    return mean_Tensor\n",
    "    \n",
    "# random logits\n",
    "def random_logits(te_scores_Tensor):\n",
    "    return te_scores_Tensor[np.random.randint(0, 1, 1)]\n",
    "\n",
    "# input: t1, t2 - triplet pair\n",
    "def triplet_distance(t1, t2):\n",
    "    return (t1 - t2).pow(2).sum()\n",
    "    \n",
    "# get triplets\n",
    "def random_triplets(st_maps, te_maps):\n",
    "    conflict = 0\n",
    "    st_triplet_list = []\n",
    "    triplet_set_size = st_maps.size(0)\n",
    "    batch_list = [x for x in range(triplet_set_size)]\n",
    "    for i in range(triplet_set_size):\n",
    "        triplet_index = random.sample(batch_list, 3)\n",
    "        anchor_index = triplet_index[0]  # denote the 1st triplet item as anchor\n",
    "        st_triplet = st_maps[triplet_index]\n",
    "        te_triplet = te_maps[triplet_index]\n",
    "        distance_01 = triplet_distance(te_triplet[0], te_triplet[1])\n",
    "        distance_02 = triplet_distance(te_triplet[0], te_triplet[2])\n",
    "        if distance_01 > distance_02:\n",
    "            conflict += 1\n",
    "            # swap postive and negative\n",
    "            st_triplet[1], st_triplet[2] = st_triplet[2], st_triplet[1]\n",
    "        st_triplet_list.append(st_triplet)\n",
    "    \n",
    "    st_triplet_batch = torch.stack(st_triplet_list, dim=1)\n",
    "    return st_triplet_batch\n",
    "    \n",
    "# get the smallest conflicts index\n",
    "def smallest_conflict_teacher(st_maps, te_maps_list):\n",
    "    \n",
    "    index = 0\n",
    "    triplet_set_size = st_maps.size(0)\n",
    "    min_conflict = 1\n",
    "    batch_list = [x for x in range(triplet_set_size)]\n",
    "    triplet_index = random.sample(batch_list, 3)\n",
    "    anchor_index = triplet_index[0]  # denote the 1st triplet item as anchor\n",
    "    for idx, te_maps in enumerate(te_maps_list):\n",
    "        conflict = 0\n",
    "        for i in range(triplet_set_size):\n",
    "            st_triplet = st_maps[triplet_index]\n",
    "            te_triplet = te_maps[triplet_index]\n",
    "            distance_01 = triplet_distance(te_triplet[0], te_triplet[1])\n",
    "            distance_02 = triplet_distance(te_triplet[0], te_triplet[2])\n",
    "            if distance_01 > distance_02:\n",
    "                conflict += 1\n",
    "        conflict /= triplet_set_size\n",
    "        conflict = min(conflict, (1-conflict))\n",
    "        if conflict < min_conflict:\n",
    "            index = idx\n",
    "    return index\n",
    "\n",
    "# train with multi-teacher\n",
    "def train(epoch, st_model):\n",
    "    print('Training:')\n",
    "    # switch to train mode\n",
    "    st_model.train()\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 = st_model(input)\n",
    "        st_maps = [b1, b2, b3, pool]\n",
    "        \n",
    "        te_scores_list = []\n",
    "        hint_maps = []\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_b2)\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",
    "        mean_logits = avg_logits(te_scores_Tensor)\n",
    "        \n",
    "        \n",
    "        te_index = smallest_conflict_teacher(b2, hint_maps)\n",
    "        st_tripets = random_triplets(b2, hint_maps[te_index])\n",
    "        \n",
    "        optimizer_sgd.zero_grad()\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        kd_loss = distillation_loss(output, target, mean_logits, T=args.T, alpha=0.7)\n",
    "        relation_loss = triplet_loss(st_tripets[0], st_tripets[1], st_tripets[2])\n",
    "        \n",
    "        loss = kd_loss + relation_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",
    "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()\n",
    "    train_loss, train_acc = train(epoch, st_model)\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",
    "        if epoch > 150:\n",
    "            torch.save(st_model.state_dict(), load_dir + st_model.model_name + '_avg.pth')\n",
    "\n",
    "# release GPU memory\n",
    "torch.cuda.empty_cache()\n",
    "log_out(\"@ BEST ACC = {:.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
}