FitNet_distill_hook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-21T02:18:39.651532Z",
     "start_time": "2019-06-21T02:18:39.184663Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Setting up a new session...\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "import os\n",
    "import time\n",
    "import logging\n",
    "import argparse\n",
    "from visdom import Visdom\n",
    "import numpy as np\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",
    "\n",
    "# Teacher models\n",
    "from models.teacher import *\n",
    "\n",
    "# Student models\n",
    "from models.student import *\n",
    "\n",
    "\n",
    "start_time = time.time()\n",
    "# os.makedirs('./checkpoint', exist_ok=True)\n",
    "\n",
    "# Training settings\n",
    "parser = argparse.ArgumentParser(description='PyTorch Distill Example')\n",
    "parser.add_argument('--teacher', type=str, default='VGG19', help='teacher net: AlexNet, VGG11/13/16/19, GoogLeNet')\n",
    "parser.add_argument('--student', type=str, default='FitNet11', help='student net: LeNet5, ')\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.01, 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('--cuda', action='store_true', default=torch.cuda.is_available(), help='use CUDA training')\n",
    "parser.add_argument('--log_interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status')\n",
    "\n",
    "config = ['--epochs', '50', '--teacher', 'Hint7', '--student', 'FitNet11', '--T', '10', '--cuda']\n",
    "args = parser.parse_args(config)\n",
    "\n",
    "device = 'cuda:0' if args.cuda and torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# logging\n",
    "logfile = './checkpoint/distill_' + args.teacher  + '_' + args.student +  '.log'\n",
    "if os.path.exists(logfile):\n",
    "    os.remove(logfile)\n",
    "\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='train 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='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",
    "# weights init\n",
    "def weights_init_normal(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    elif classname.find(\"BatchNorm2d\") != -1:\n",
    "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        nn.init.constant_(m.bias.data, 0.0)\n",
    "    elif classname.find('linear') != -1:\n",
    "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        nn.init.constant_(m.bias.data, 0.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class ConvRegressor(nn.Module):\n",
    "#     def __init__(self, teacher, hint_layer, student, guided_layer):\n",
    "#         self.hint_layer = teacher.\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-21T08:25:23.258466Z",
     "start_time": "2019-06-21T08:25:22.836761Z"
    },
    "code_folding": []
   },
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: './checkpoint/Hint7_cifar10.pth'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-55f3d95d3149>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mteacher_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mteacher\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mteacher_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload_state_dict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'./checkpoint/'\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mteacher\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'_cifar10.pth'\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      3\u001b[0m \u001b[0mst_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstudent\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mst_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweights_init_normal\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# init student\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Dev/anaconda3/envs/py37/lib/python3.7/site-packages/torch/serialization.py\u001b[0m in \u001b[0;36mload\u001b[0;34m(f, map_location, pickle_module, **pickle_load_args)\u001b[0m\n\u001b[1;32m    380\u001b[0m             \u001b[0;34m(\u001b[0m\u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mversion_info\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m2\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0municode\u001b[0m\u001b[0;34m)\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    381\u001b[0m         \u001b[0mnew_fd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 382\u001b[0;31m         \u001b[0mf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\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    383\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mversion_info\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m3\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpathlib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m)\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    384\u001b[0m         \u001b[0mnew_fd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: './checkpoint/Hint7_cifar10.pth'"
     ]
    }
   ],
   "source": [
    "teacher_model = eval(args.teacher)().to(device)\n",
    "teacher_model.load_state_dict(torch.load('./checkpoint/' + args.teacher + '_cifar10.pth'))\n",
    "st_model = eval(args.student)().to(device)\n",
    "st_model.apply(weights_init_normal)  # init student\n",
    "\n",
    "st_features = None\n",
    "te_features = None\n",
    "\n",
    "def st_hook(module, input, output):\n",
    "    '''把这层的输出到features中'''\n",
    "    global st_features\n",
    "    st_features = output.data\n",
    "    \n",
    "def te_hook(module, input, output):\n",
    "    '''把这层的输出拷贝到features中'''\n",
    "    global te_features\n",
    "    te_features = output.data\n",
    "\n",
    "class Regressor(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Regressor,self).__init__()\n",
    "        # torch.Size([128, 512, 16, 16]) -> torch.Size([128, 80, 4, 4]) \n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(2048, 512, 3, 1, 1),  # ch: 512 -> 256  \n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2), # size: 16 -> 8\n",
    "            \n",
    "            nn.Conv2d(512, 128, 3, 1, 1),  # ch: 256 -> 128\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2), # size: 8 -> 4\n",
    "            \n",
    "            nn.Conv2d(128, 80, 3, 1, 1),   # ch: 128 -> 80\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.features(x)\n",
    "\n",
    "regressor = Regressor().to(device)\n",
    "st_model.features[15].register_forward_hook(st_hook)\n",
    "teacher_model.features[11].register_forward_hook(te_hook)\n",
    "\n",
    "# data\n",
    "train_transform = transforms.Compose([transforms.RandomHorizontalFlip(), transforms.ToTensor()])\n",
    "test_transform = transforms.Compose([transforms.ToTensor()])\n",
    "train_set = datasets.CIFAR10(root='../data', train=True, download=True, transform=train_transform)\n",
    "test_set = datasets.CIFAR10(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_r = optim.SGD(regressor.parameters(), lr=args.lr, momentum=args.momentum)\n",
    "\n",
    "def distillation(y, labels, teacher_scores, T, alpha):\n",
    "    return nn.KLDivLoss()(F.log_softmax(y/T), F.softmax(teacher_scores/T)) * (T*T * 2.0 * alpha) + F.cross_entropy(y, labels) * (1. - alpha)\n",
    "\n",
    "\n",
    "# guided train\n",
    "def guided_train(model, loss_fn_guided=nn.MSELoss()):\n",
    "    model.train()\n",
    "    teacher_model.eval()\n",
    "    guided_loss = None\n",
    "    for epoch in range(10):\n",
    "        for batch_idx, (data, target) in enumerate(train_loader):\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            optimizer_r.zero_grad()\n",
    "            model.forward(data)\n",
    "            teacher_model(data).detach()\n",
    "            te_output = regressor(te_features)\n",
    "#             print(st_features.size(), te_hint.size())\n",
    "        \n",
    "            guided_loss = loss_fn_guided(te_output, st_features)\n",
    "            guided_loss.backward()\n",
    "            optimizer_r.step()\n",
    "            optimizer.step()\n",
    "        print('guided_epoch:[{}]\\tLoss:{:.4f}'.format(epoch, guided_loss.item()))\n",
    "\n",
    "\n",
    "def train(epoch, model, loss_fn):\n",
    "    model.train()\n",
    "    teacher_model.eval()\n",
    "    loss = None\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        teacher_output = teacher_model(data)\n",
    "        teacher_output = teacher_output.detach()\n",
    "#         print(st_features.size())\n",
    "        # teacher_output = Variable(teacher_output.data, requires_grad=False) #alternative approach to load teacher_output\n",
    "        loss = loss_fn(output, target, teacher_output, T=args.T, alpha=0.6)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx % args.log_interval == 0:\n",
    "            log_out('Train Epoch: {} [{}/{} ({:.4f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                100. * batch_idx / len(train_loader), loss.item()))\n",
    "    return loss.item()\n",
    "\n",
    "def train_evaluate(model):\n",
    "    model.eval()\n",
    "    train_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in train_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            train_loss += F.cross_entropy(output, target).item() # sum up batch loss\n",
    "            pred = output.data.max(1, keepdim=True)[1]\n",
    "            correct += pred.eq(target.data.view_as(pred)).cpu().sum()\n",
    "\n",
    "        log_out('\\nTrain set: Average loss: {:.4f}, Accuracy: {}/{} ({:.4f}%)\\n'.format(\n",
    "            train_loss, correct, len(train_loader.dataset), \n",
    "            100. * correct / len(train_loader.dataset)))\n",
    "    return 100. * correct / len(train_loader.dataset)\n",
    "\n",
    "def test(model):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            # test_loss += F.cross_entropy(output, target).item() # sum up batch loss\n",
    "            pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability\n",
    "            correct += pred.eq(target.data.view_as(pred)).cpu().sum()\n",
    "\n",
    "        test_loss /= len(test_loader.dataset)\n",
    "        log_out('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.4f}%)\\n'.format(\n",
    "            test_loss, correct, len(test_loader.dataset),\n",
    "            100. * correct / len(test_loader.dataset)))\n",
    "    return 100. * correct / len(test_loader.dataset)\n",
    "\n",
    "print('StudentNet:\\n')\n",
    "print(st_model)\n",
    "guided_train(st_model)\n",
    "for epoch in range(1, args.epochs + 1):\n",
    "    train_loss = train(epoch, st_model, loss_fn=distillation)\n",
    "    # visaulize loss\n",
    "    vis.line(np.array([train_loss]), np.array([epoch]), loss_win, update=\"append\")\n",
    "    train_acc = train_evaluate(st_model)\n",
    "    test_acc = test(st_model)\n",
    "    vis.line(np.column_stack((train_acc, test_acc)), np.column_stack((epoch, epoch)), acc_win, update=\"append\")\n",
    "\n",
    "\n",
    "torch.save(st_model.state_dict(), './checkpoint/' + args.teacher + '_distill_' + args.student + '.pth')\n",
    "# the_model = Net()\n",
    "# the_model.load_state_dict(torch.load('student.pth.tar'))\n",
    "\n",
    "# test(the_model)\n",
    "# for data, target in test_loader:\n",
    "#     data, target = Variable(data, volatile=True), Variable(target)\n",
    "#     teacher_out = the_model(data)\n",
    "# print(teacher_out)\n",
    "log_out(\"--- {:.3f} seconds ---\".format(time.time() - start_time))\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
}