meta-quantize-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is meant for providing a detailed description of MetaQuant via code and mathematical explanation.\n",
    "\n",
    "## Motivation\n",
    "Training-based quantization aims at minimizing the following training loss:\n",
    "\n",
    "$$\\min \\ell = \\text{Loss}(f(Q(\\mathbf{W}, \\mathbf{x})))$$\n",
    "\n",
    "where $Q(\\cdot)$ quantize full-precision weight $\\mathbf{W}$ into quantized value $\\mathbf{\\hat{W}}$. Due to the non-differentiability of $Q(\\cdot)$, the gradient of $\\ell$ w.r.t $\\mathbf{W}$ cannot be attained in a normal way.\n",
    "To enable a stable quantization training, Straight-Through-Estimator (STE) is proposed to redefine $\\partial Q(r)/\\partial r$:\n",
    "\n",
    "\\begin{eqnarray}\n",
    "\t\\frac{\\partial Q(r)}{\\partial r} =\\left\\{\n",
    "\t\\begin{aligned}\n",
    "\t1 & \\qquad \\text{if} \\qquad |r| \\leq 1  \\nonumber \\\\\n",
    "\t0 & \\qquad \\text{else}.\n",
    "\t\\end{aligned}\n",
    "\t\\right..\n",
    "\\end{eqnarray}\n",
    "\n",
    "However, it inevitably brings the problem of **gradient mismatch**: the gradients of the weights are not generated using the value of weights, but rather its quantized value. Although STE provides an end-to-end training method under discrete constraints, few works have progressed to investigate how to obtain better gradients for quantization training. \n",
    "\n",
    "To overcome the problem of gradient mismatch and explore better gradients in training-based methods, we propose to learn $\\frac{\\partial Q(\\mathbf{W})}{\\partial \\mathbf{W}}$ by a neural network ($\\mathcal{W}$) during quantization training. Such neural network is called **meta quantizer** and is trained together with the base quantized model. This process is named as **Meta** **Quant**ization (MetaQuant). \n",
    "\n",
    "Specially, in each backward propagation, $\\mathcal{W}$ takes $\\frac{\\partial \\ell}{\\partial Q(\\mathbf{W})}$ and $\\mathbf{W}$ as inputs in a coordinate-wise manner, then its output is assigned to $\\frac{\\partial \\ell}{\\partial Q(\\mathbf{W})}$ for weights update using common optimization methods such as SGD and Adam. In the forward pass, inference is conducted using the quantized version of the updated weights, which produce the final outputs to be compared with the ground-truth labels for backward computation. During this process, gradient propagation from the quantized weights to the full-precision weights is handled by $\\mathcal{M}$, which avoids the problem of non-differentiability and gradient mismatch. Besides, the gradients generated by the \\meta are loss-aware, contributing to better performance of the quantization training.\n",
    "\n",
    "## Overflow of MetaQuant\n",
    "<!--\n",
    "![Overflow of MetaQuant]('./figs/MetaQuant.png')\n",
    "-->\n",
    "\n",
    "<img src=\"figs/MetaQuant.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This block import packages and initialize key parameters\n",
    "\"\"\"\n",
    "import os\n",
    "# os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "import shutil\n",
    "import pickle\n",
    "import time\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim as optim\n",
    "\n",
    "from utils.dataset import get_dataloader\n",
    "from meta_utils.meta_network import MetaFC, MetaLSTMFC, MetaDesignedMultiFC\n",
    "from meta_utils.SGD import SGD\n",
    "from meta_utils.adam import Adam\n",
    "from meta_utils.helpers import meta_gradient_generation, update_parameters\n",
    "from utils.recorder import Recorder\n",
    "from utils.miscellaneous import AverageMeter, accuracy, progress_bar\n",
    "from utils.miscellaneous import get_layer\n",
    "from utils.quantize import test\n",
    "\n",
    "from models_CIFAR.quantized_meta_resnet import resnet20_cifar\n",
    "\n",
    "# ------------------------------------------\n",
    "use_cuda = torch.cuda.is_available()\n",
    "model_name = 'ResNet20'\n",
    "dataset_name = 'CIFAR10'\n",
    "meta_method = 'MultiFC' # ['LSTMFC', 'MultiFC', 'FC-Grad']\n",
    "MAX_EPOCH = 100\n",
    "optimizer_type = 'SGD' # ['SGD', 'SGD-M', 'adam']\n",
    "hidden_size = 100\n",
    "num_lstm = 2\n",
    "num_fc = 3\n",
    "lr_adjust = '30'\n",
    "batch_size = 128\n",
    "bitW = 1\n",
    "quantized_type = 'dorefa' # ['dorefa', 'BWN']\n",
    "save_root = './Results/%s-%s' % (model_name, dataset_name)\n",
    "# ------------------------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This block initialize network and load dataset\n",
    "\"\"\"\n",
    "\n",
    "import utils.global_var as gVar\n",
    "gVar.meta_count = 0\n",
    "\n",
    "###################\n",
    "# Initial Network #\n",
    "###################\n",
    "net = resnet20_cifar(bitW=bitW)\n",
    "pretrain_path = '%s/%s-%s-pretrain.pth' % (save_root, model_name, dataset_name)\n",
    "net.load_state_dict(torch.load(pretrain_path), strict=False)\n",
    "\n",
    "# Get layer name list\n",
    "layer_name_list = net.layer_name_list\n",
    "# Assert all required layer is initialized as meta layer\n",
    "assert (len(layer_name_list) == gVar.meta_count)\n",
    "print('Layer name list completed.')\n",
    "\n",
    "if use_cuda:\n",
    "    net.cuda()\n",
    "    \n",
    "################\n",
    "# Load Dataset #\n",
    "################\n",
    "train_loader = get_dataloader(dataset_name, 'train', batch_size)\n",
    "test_loader = get_dataloader(dataset_name, 'test', 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This block initialize meta network, optimizer and recorder\n",
    "\"\"\"\n",
    "########################\n",
    "# Initial Meta Network #\n",
    "########################\n",
    "if meta_method == 'LSTMFC':\n",
    "    meta_net = MetaLSTMFC(hidden_size=hidden_size)\n",
    "    SummaryPath = '%s/runs-Quant/Meta-%s-Nonlinear-%s-' \\\n",
    "                  'hidden-size-%d-nlstm-1-%s-%s-%dbits-lr-%s' \\\n",
    "                  % (save_root, meta_method, args.meta_nonlinear, hidden_size,\n",
    "                     quantized_type, optimizer_type, bitW, lr_adjust)\n",
    "elif meta_method in ['FC-Grad']:\n",
    "    meta_net = MetaFC(hidden_size=hidden_size, use_nonlinear=args.meta_nonlinear)\n",
    "    SummaryPath = '%s/runs-Quant/Meta-%s-Nonlinear-%s-' \\\n",
    "                  'hidden-size-%d-%s-%s-%dbits-lr-%s' \\\n",
    "                  % (save_root, meta_method, args.meta_nonlinear, hidden_size,\n",
    "                     quantized_type, optimizer_type, bitW, lr_adjust)\n",
    "elif meta_method == 'MultiFC':\n",
    "    meta_net = MetaDesignedMultiFC(hidden_size=hidden_size,\n",
    "                                   num_layers = args.num_fc,\n",
    "                                   use_nonlinear=args.meta_nonlinear)\n",
    "    SummaryPath = '%s/runs-Quant/Meta-%s-Nonlinear-%s-' \\\n",
    "                  'hidden-size-%d-nfc-%d-%s-%s-%dbits-lr-%s' \\\n",
    "                  % (save_root, meta_method, args.meta_nonlinear, hidden_size, num_fc,\n",
    "                     quantized_type, optimizer_type, bitW, lr_adjust)\n",
    "else:\n",
    "    raise NotImplementedError\n",
    "\n",
    "print(meta_net)\n",
    "\n",
    "if use_cuda:\n",
    "    meta_net.cuda()\n",
    "\n",
    "meta_optimizer = optim.Adam(meta_net.parameters(), lr=1e-3, weight_decay=args.weight_decay)\n",
    "    \n",
    "#####################\n",
    "# Initial Optimizee #\n",
    "#####################\n",
    "    \n",
    "# Optimizer for original network, just for zeroing gradient and get refined gradient\n",
    "if optimizer_type == 'SGD-M':\n",
    "    optimizee = SGD(net.parameters(), lr=args.init_lr,\n",
    "                    momentum=0.9, weight_decay=5e-4)\n",
    "elif optimizer_type == 'SGD':\n",
    "    optimizee = SGD(net.parameters(), lr=args.init_lr)\n",
    "elif optimizer_type in ['adam', 'Adam']:\n",
    "    optimizee = Adam(net.parameters(), lr=args.init_lr,\n",
    "                     weight_decay=5e-4)\n",
    "else:\n",
    "    raise NotImplementedError\n",
    "    \n",
    "####################\n",
    "# Initial Recorder #\n",
    "####################\n",
    "if args.exp_spec is not '':\n",
    "    SummaryPath += ('-' + args.exp_spec)\n",
    "\n",
    "print('Save to %s' %SummaryPath)\n",
    "\n",
    "if os.path.exists(SummaryPath):\n",
    "    print('Record exist, remove')\n",
    "    input()\n",
    "    shutil.rmtree(SummaryPath)\n",
    "    os.makedirs(SummaryPath)\n",
    "else:\n",
    "    os.makedirs(SummaryPath)\n",
    "\n",
    "recorder = Recorder(SummaryPath=SummaryPath, dataset_name=dataset_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In training of MetaQuant, in order to train meta quantizer, the meta gradient produced by meta quantizer is added to the inference process to get loss:\n",
    "\n",
    "<img src=\"figs/MetaQuant-Forward.png\">\n",
    "\n",
    "Therefore, the forward process in base net is modified to incorprate meta gradient embedded using ```meta_grad_dict```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This block begins training\n",
    "\"\"\"\n",
    "meta_hidden_state_dict = dict() # Dictionary to store hidden states for all layers for memory-based meta network\n",
    "meta_grad_dict = dict() # Dictionary to store meta net output: gradient for origin network's weight / bias\n",
    "\n",
    "for epoch in range(MAX_EPOCH):\n",
    "\n",
    "    if recorder.stop: break\n",
    "\n",
    "    print('\\nEpoch: %d, lr: %e' % (epoch, optimizee.param_groups[0]['lr']))\n",
    "\n",
    "    net.train()\n",
    "    end = time.time()\n",
    "\n",
    "    recorder.reset_performance()\n",
    "\n",
    "    for batch_idx, (inputs, targets) in enumerate(train_loader):\n",
    "\n",
    "        if use_cuda:\n",
    "            inputs, targets = inputs.cuda(), targets.cuda()\n",
    "\n",
    "        meta_optimizer.zero_grad()\n",
    "\n",
    "        # In first iteration of whole training, meta gradient hasn't been generated, \n",
    "        # therefore the first forward is conducted without meta gradient.\n",
    "        if batch_idx == 0 and epoch == 0:\n",
    "            pass\n",
    "        # meta gradient used in current iteration is generated by the gradient and weights \n",
    "        # from previous iteration.\n",
    "        else:\n",
    "            meta_grad_dict, meta_hidden_state_dict = \\\n",
    "                meta_gradient_generation(\n",
    "                        meta_net, net, meta_method, meta_hidden_state_dict\n",
    "                )\n",
    "        # Conduct forward using meta gradient\n",
    "        outputs = net(inputs, quantized_type=quantized_type,\n",
    "                      meta_grad_dict=meta_grad_dict,\n",
    "                      lr=optimizee.param_groups[0]['lr'])\n",
    "\n",
    "        optimizee.zero_grad()\n",
    "\n",
    "        # Taking backward generate gradient for meta pruner and base model\n",
    "        # Non-meta-weights' (bias, BN layer) gradient is attained here\n",
    "        losses = nn.CrossEntropyLoss()(outputs, targets)\n",
    "        losses.backward()\n",
    "\n",
    "        meta_optimizer.step()\n",
    "\n",
    "        # Assign meta gradient for actual gradients used in update_parameters\n",
    "        if len(meta_grad_dict) != 0:\n",
    "            for layer_info in net.layer_name_list:\n",
    "                layer_name = layer_info[0]\n",
    "                layer_idx = layer_info[1]\n",
    "                layer = get_layer(net, layer_idx)\n",
    "                layer.weight.grad.data = (layer.calibration * layer.pre_quantized_grads)\n",
    "                # layer.weight.grad.data.copy_(layer.calibration * meta_grad_dict[layer_name][1].data)\n",
    "\n",
    "        # Get refine gradients for next computation\n",
    "        optimizee.get_refine_gradient()\n",
    "\n",
    "        # These gradient should be saved in next iteration's inference\n",
    "        if len(meta_grad_dict) != 0:\n",
    "            update_parameters(net, lr=optimizee.param_groups[0]['lr'])\n",
    "\n",
    "        recorder.update(loss=losses.data.item(), acc=accuracy(outputs.data, targets.data, (1,5)),\n",
    "                        batch_size=outputs.shape[0], cur_lr=optimizee.param_groups[0]['lr'], end=end)\n",
    "\n",
    "        recorder.print_training_result(batch_idx, len(train_loader))\n",
    "        end = time.time()\n",
    "\n",
    "    test_acc = test(net, quantized_type=quantized_type, test_loader=test_loader,\n",
    "                    dataset_name=dataset_name, n_batches_used=None)\n",
    "    recorder.update(loss=None, acc=test_acc, batch_size=0, end=None, is_train=False)\n",
    "\n",
    "    # Adjust learning rate\n",
    "    recorder.adjust_lr(optimizer=optimizee, adjust_type=lr_adjust, epoch=epoch)\n",
    "\n",
    "best_test_acc = recorder.get_best_test_acc()\n",
    "if type(best_test_acc) == tuple:\n",
    "    print('Best test top 1 acc: %.3f, top 5 acc: %.3f' % (best_test_acc[0], best_test_acc[1]))\n",
    "else:\n",
    "    print('Best test acc: %.3f' %best_test_acc)\n",
    "recorder.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}