Resnet_18_Imagenet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compressing a Resnet-18 trained on Imagenet\n",
    "\n",
    "#### To replicate results with pretrained models please download the following models\n",
    "\n",
    "1. resnet18_imagenet_pytorch.pth\n",
    "2. resnet18_imagenet_pytorch_small_cup_t_point_9.pth\n",
    "3. resnet18_imagenet_pytorch_lr_point1_cupSS_K_point03_b_point3.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys; sys.argv=['']; \n",
    "sys.path.insert(0, '../')\n",
    "del sys\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '0,1,2'\n",
    "\n",
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.nn.parallel\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.distributed as dist\n",
    "import torch.multiprocessing as mp\n",
    "import torch.utils.data\n",
    "import torch.utils.data.distributed\n",
    "import torch.utils.model_zoo as model_zoo\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "from torchvision import datasets, transforms\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "import time\n",
    "import copy\n",
    "\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = '1'\n",
    "\n",
    "from src.imagenet_utils import train,validate,save_checkpoint,AverageMeter,ProgressMeter\n",
    "from src.imagenet_utils import adjust_learning_rate,accuracy,adjust_learning_rate_pytorch_retrain\n",
    "from src.utils import plot_tsne,fancy_dendrogram,save_obj,load_obj,weights_init\n",
    "from src.model import VGG,load_model\n",
    "from src.prune_model import prune_model\n",
    "from src.cluster_model import cluster_model\n",
    "from src.train_test import adjust_learning_rate_nips,adjust_learning_rate_iccv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specify imagenet data path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "_StoreAction(option_strings=['--data'], dest='data', nargs=None, const=None, default='/localscratch/mgh/Rahul_Imagenet/', type=<class 'str'>, choices=None, help='path to dataset', metavar='S')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')\n",
    "\n",
    "### add path to dataset here #####\n",
    "parser.add_argument('--data', type=str,default='/localscratch/mgh/Rahul_Imagenet/',metavar='S',help='path to dataset')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "parser.add_argument('--start-epoch', default=0, type=int, metavar='N',\n",
    "                    help='manual epoch number (useful on restarts)')\n",
    "parser.add_argument('-b', '--batch-size', default=256, type=int,\n",
    "                    metavar='N',\n",
    "                    help='mini-batch size (default: 256), this is the total '\n",
    "                         'batch size of all GPUs on the current node when '\n",
    "                         'using Data Parallel or Distributed Data Parallel')\n",
    "parser.add_argument('-j', '--workers', default=6, type=int, metavar='N',\n",
    "                    help='number of data loading workers (default: 4)')\n",
    "parser.add_argument('--epochs', default=90, type=int, metavar='N',\n",
    "                    help='number of total epochs to run')\n",
    "parser.add_argument('--lr', '--learning-rate', default=0.01, type=float,\n",
    "                    metavar='LR', help='initial learning rate')\n",
    "parser.add_argument('--momentum', default=0.9, type=float, metavar='M',\n",
    "                    help='momentum')\n",
    "parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,\n",
    "                    metavar='W', help='weight decay (default: 1e-4)')\n",
    "parser.add_argument('-p', '--print-freq', default=500, type=int,\n",
    "                    metavar='N', help='print frequency (default: 10)')\n",
    "parser.add_argument('--resume', default='', type=str, metavar='PATH',\n",
    "                    help='path to latest checkpoint (default: none)')\n",
    "parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',\n",
    "                    help='evaluate model on validation set')\n",
    "parser.add_argument('--pretrained', dest='pretrained', action='store_true',\n",
    "                    help='use pre-trained model')\n",
    "parser.add_argument('--world-size', default=-1, type=int,\n",
    "                    help='number of nodes for distributed training')\n",
    "parser.add_argument('--rank', default=-1, type=int,\n",
    "                    help='node rank for distributed training')\n",
    "parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,\n",
    "                    help='url used to set up distributed training')\n",
    "parser.add_argument('--dist-backend', default='nccl', type=str,\n",
    "                    help='distributed backend')\n",
    "parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "                    help='disables CUDA training')\n",
    "parser.add_argument('--seed', type=int, default=12346, metavar='S',\n",
    "                    help='random seed (default: 12346)')\n",
    "parser.add_argument('--num_output', type=int, default=10, metavar='S',\n",
    "                    help='number of classes(default: 10)')\n",
    "parser.add_argument('--log-interval', type=int, default=100, metavar='N',\n",
    "                    help='how many batches to wait before logging training status')\n",
    "parser.add_argument('--checkpoint_path', type=str, default='./checkpoints/resnet18_imagenet_pytorch.pth', metavar='S',\n",
    "                    help='path to store model training checkpoints')\n",
    "parser.add_argument('--gpu', type=int, default=0, nargs='+', help='used gpu')\n",
    "parser.add_argument('--multiprocessing-distributed', action='store_true',\n",
    "                        help='Use multi-processing distributed training to launch ')#,\n",
    "#                          'N processes per node, which has N GPUs. This is the ',\n",
    "#                          'fastest way to use PyTorch for either single node or ',\n",
    "#                          'multi node data parallel training')\n",
    "\n",
    "args = parser.parse_args()\n",
    "\n",
    "use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "# if use_cuda:\n",
    "#     print('using gpu',args.gpu)\n",
    "#     os.environ[\"CUDA_VISIBLE_DEVICES\"] = ','.join(str(x) for x in args.gpu)\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "\n",
    "# Data loading code\n",
    "traindir = os.path.join(args.data, 'train')\n",
    "valdir = os.path.join(args.data, 'val')\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "\n",
    "train_dataset = datasets.ImageFolder(\n",
    "    traindir,\n",
    "    transforms.Compose([\n",
    "        transforms.RandomResizedCrop(224),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        normalize,\n",
    "    ]))\n",
    "\n",
    "train_sampler = None\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),\n",
    "    num_workers=args.workers, pin_memory=True, sampler=train_sampler)\n",
    "\n",
    "val_loader = torch.utils.data.DataLoader(\n",
    "    datasets.ImageFolder(valdir, transforms.Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        normalize,\n",
    "    ])),\n",
    "    batch_size=args.batch_size, shuffle=False,\n",
    "    num_workers=args.workers, pin_memory=True)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss().cuda(args.gpu)\n",
    "\n",
    "\n",
    "writer = SummaryWriter('logs/resnet18_imagenet/')\n",
    "\n",
    "#set all seeds for reproducability\n",
    "def set_random_seed(seed):    \n",
    "    random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    np.random.seed(args.seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "set_random_seed(args.seed)\n",
    "\n",
    "torch.cuda.set_device(args.gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train baseline model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " * Acc@1 69.876 Acc@5 89.256\n",
      "loaded model with top1 : 69.8759994506836, top5 : 89.25599670410156\n"
     ]
    }
   ],
   "source": [
    "set_random_seed(args.seed)\n",
    "\n",
    "args.lr = 0.1\n",
    "\n",
    "resnet18 = torchvision.models.resnet18(pretrained=False)\n",
    "optimizer = torch.optim.SGD(resnet18.parameters(), args.lr,\n",
    "                            momentum=args.momentum,\n",
    "                            weight_decay=args.weight_decay)\n",
    "\n",
    "resnet18 = torch.nn.DataParallel(resnet18)\n",
    "resnet18.cuda(args.gpu)\n",
    "\n",
    "best_val_acc = 0\n",
    "\n",
    "if not os.path.isfile(args.checkpoint_path):\n",
    "    for epoch in range(args.start_epoch, args.epochs):\n",
    "        adjust_learning_rate_pytorch_retrain(optimizer, epoch, args)        \n",
    "\n",
    "        # train for one epoch\n",
    "        train_loss,train_top1,train_top5 = train(train_loader, resnet18, criterion, optimizer, epoch, args)\n",
    "\n",
    "        # evaluate on validation set\n",
    "        val_loss,val_top1,val_top5 = validate(val_loader, resnet18, criterion, args)\n",
    "        \n",
    "        if val_top1 > best_val_acc:  \n",
    "            torch.save(resnet18, args.checkpoint_path, pickle_protocol=4)            \n",
    "            best_val_acc = val_top1    \n",
    "\n",
    "        writer.add_scalars('resnet18_imagenet_pytorch_schedule/loss',{'train_loss': train_loss,\n",
    "                                        'val_loss' : val_loss}, epoch)\n",
    "        writer.add_scalars('resnet18_imagenet_pytorch_schedule/accuracy',{'train_top1': train_top1,\n",
    "                                                  'val_top1': val_top1,\n",
    "                                                  'train_top5': train_top5,\n",
    "                                                  'val_top5': val_top5}, epoch) \n",
    "else:   \n",
    "    resnet18 = torch.load(args.checkpoint_path)\n",
    "    _,val_top1,val_top5 = validate(val_loader, resnet18, criterion, args, verbose=False)\n",
    "    \n",
    "print('loaded model with top1 : {}, top5 : {}'.format(val_top1,val_top5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compress using CUP (T = 0.8)\n",
    "\n",
    "- This section compresses the resnet model that we trained"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load pre-trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " * Acc@1 69.876 Acc@5 89.256\n",
      "loaded model with top1 : 69.8759994506836, top5 : 89.25599670410156\n"
     ]
    }
   ],
   "source": [
    "resnet18 = torch.load(args.checkpoint_path).module\n",
    "_,val_top1,val_top5 = validate(val_loader, resnet18, criterion, args, verbose=False)\n",
    "best_val_acc = val_top5\n",
    "print('loaded model with top1 : {}, top5 : {}'.format(val_top1,val_top5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " * Acc@1 0.294 Acc@5 0.938\n",
      " * Acc@1 68.876 Acc@5 88.466\n",
      "large top-1 69.88, small top-1 68.88, top-1 drop 1.00, large top-5 89.26, small top-5 88.47, top-5 drop 0.79\n"
     ]
    }
   ],
   "source": [
    "cluster_args = {    \n",
    "    'cluster_layers' : {4:0,9:0,14:0,21:0,26:0,33:0,38:0,45:0},\n",
    "    'conv_feature_size' : 1,\n",
    "    'features' : 'both',\n",
    "    'channel_reduction' : 'fro',\n",
    "    'use_bias' : False,\n",
    "    'reshape_exists' : False,\n",
    "    'linkage_method' : 'ward',\n",
    "    'distance_metric' : 'euclidean',\n",
    "    'cluster_criterion' : 'hierarchical',\n",
    "    'distance_threshold' : 0.8,\n",
    "    'merge_criterion' : 'max_l2_norm',\n",
    "    'verbose' : False\n",
    "}\n",
    "\n",
    "path = args.checkpoint_path[:-4] + '_small_cup_t_point_8.pth' \n",
    "model_modifier = cluster_model(resnet18,cluster_args)\n",
    "resnet18_clustered = model_modifier.cluster_model()\n",
    "optimizer = torch.optim.SGD(resnet18_clustered.parameters(), args.lr,\n",
    "                            momentum=args.momentum,\n",
    "                            weight_decay=args.weight_decay)\n",
    "\n",
    "resnet18_clustered = torch.nn.DataParallel(resnet18_clustered)\n",
    "resnet18_clustered.cuda(args.gpu)\n",
    "\n",
    "_,top1_acc_no_retrain,top5_acc_no_retrain = validate(val_loader, resnet18_clustered, criterion, args, verbose=False)\n",
    "\n",
    "best_val_acc = 0\n",
    "\n",
    "if not os.path.isfile(path):\n",
    "    for epoch in range(args.start_epoch, args.epochs):\n",
    "        adjust_learning_rate_pytorch_retrain(optimizer, epoch, args)        \n",
    "\n",
    "        # train for one epoch\n",
    "        train_loss,train_top1,train_top5 = train(train_loader, resnet18_clustered, criterion, optimizer, epoch, args)\n",
    "\n",
    "        # evaluate on validation set\n",
    "        val_loss,val_top1,val_top5 = validate(val_loader, resnet18_clustered, criterion, args)\n",
    "        \n",
    "        if val_top1 > best_val_acc:  \n",
    "            torch.save(resnet18_clustered, path, pickle_protocol=4)            \n",
    "            best_val_acc = val_top1    \n",
    "\n",
    "        writer.add_scalars('resnet18_imagenet_t_point_8_pytorch/loss',{'train_loss': train_loss,\n",
    "                                        'val_loss' : val_loss}, epoch)\n",
    "        writer.add_scalars('resnet18_imagenet_t_point_8_pytorch/accuracy',{'train_top1': train_top1,\n",
    "                                                  'val_top1': val_top1,\n",
    "                                                  'train_top5': train_top5,\n",
    "                                                  'val_top5': val_top5}, epoch) \n",
    "else:   \n",
    "    resnet18_clustered = torch.load(path)\n",
    "    _,test_top1_acc,test_top5_acc = validate(val_loader, resnet18_clustered, criterion, args, verbose=False)\n",
    "        \n",
    "print('large top-1 {:.2f}, small top-1 {:.2f}, top-1 drop {:.2f}, large top-5 {:.2f}, small top-5 {:.2f}, top-5 drop {:.2f}'.format(val_top1,test_top1_acc,val_top1-test_top1_acc,val_top5,test_top5_acc,val_top5-test_top5_acc))       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  + Number of FLOPs: 3.64G\n",
      "  + Number of FLOPs: 2.08G\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2077024562.0"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from src.compute_flops import print_model_param_nums,print_model_param_flops\n",
    "\n",
    "print_model_param_flops(resnet18.cpu(),input_res=224)\n",
    "print_model_param_flops(resnet18_clustered.module.cpu(),input_res=224)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compress using CUP (T = 0.825)\n",
    "\n",
    "- This section compresses the resnet model that we trained"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load pre-trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " * Acc@1 69.876 Acc@5 89.256\n",
      "loaded model with top1 : 69.8759994506836, top5 : 89.25599670410156\n"
     ]
    }
   ],
   "source": [
    "resnet18 = torch.load(args.checkpoint_path).module\n",
    "_,val_top1,val_top5 = validate(val_loader, resnet18, criterion, args, verbose=False)\n",
    "best_val_acc = val_top5\n",
    "print('loaded model with top1 : {}, top5 : {}'.format(val_top1,val_top5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " * Acc@1 0.284 Acc@5 0.916\n",
      " * Acc@1 68.292 Acc@5 88.146\n",
      "large top-1 69.88, small top-1 68.29, top-1 drop 1.58, large top-5 89.26, small top-5 88.15, top-5 drop 1.11\n"
     ]
    }
   ],
   "source": [
    "cluster_args = {    \n",
    "    'cluster_layers' : {4:0,9:0,14:0,21:0,26:0,33:0,38:0,45:0},\n",
    "    'conv_feature_size' : 1,\n",
    "    'features' : 'both',\n",
    "    'channel_reduction' : 'fro',\n",
    "    'use_bias' : False,\n",
    "    'reshape_exists' : False,\n",
    "    'linkage_method' : 'ward',\n",
    "    'distance_metric' : 'euclidean',\n",
    "    'cluster_criterion' : 'hierarchical',\n",
    "    'distance_threshold' : 0.825,\n",
    "    'merge_criterion' : 'max_l2_norm',\n",
    "    'verbose' : False\n",
    "}\n",
    "\n",
    "path = args.checkpoint_path[:-4] + '_small_cup_t_point_825.pth' \n",
    "model_modifier = cluster_model(resnet18,cluster_args)\n",
    "resnet18_clustered = model_modifier.cluster_model()\n",
    "optimizer = torch.optim.SGD(resnet18_clustered.parameters(), args.lr,\n",
    "                            momentum=args.momentum,\n",
    "                            weight_decay=args.weight_decay)\n",
    "\n",
    "resnet18_clustered = torch.nn.DataParallel(resnet18_clustered)\n",
    "resnet18_clustered.cuda(args.gpu)\n",
    "\n",
    "_,top1_acc_no_retrain,top5_acc_no_retrain = validate(val_loader, resnet18_clustered, criterion, args, verbose=False)\n",
    "\n",
    "best_val_acc = 0\n",
    "\n",
    "if not os.path.isfile(path):\n",
    "    for epoch in range(args.start_epoch, args.epochs):\n",
    "        adjust_learning_rate_pytorch_retrain(optimizer, epoch, args)        \n",
    "\n",
    "        # train for one epoch\n",
    "        train_loss,train_top1,train_top5 = train(train_loader, resnet18_clustered, criterion, optimizer, epoch, args)\n",
    "\n",
    "        # evaluate on validation set\n",
    "        val_loss,val_top1,val_top5 = validate(val_loader, resnet18_clustered, criterion, args)\n",
    "        \n",
    "        if val_top1 > best_val_acc:  \n",
    "            torch.save(resnet18_clustered, path, pickle_protocol=4)            \n",
    "            best_val_acc = val_top1    \n",
    "\n",
    "        writer.add_scalars('resnet18_imagenet_t_point_825_pytorch/loss',{'train_loss': train_loss,\n",
    "                                        'val_loss' : val_loss}, epoch)\n",
    "        writer.add_scalars('resnet18_imagenet_t_point_825_pytorch/accuracy',{'train_top1': train_top1,\n",
    "                                                  'val_top1': val_top1,\n",
    "                                                  'train_top5': train_top5,\n",
    "                                                  'val_top5': val_top5}, epoch) \n",
    "else:   \n",
    "    resnet18_clustered = torch.load(path)\n",
    "    _,test_top1_acc,test_top5_acc = validate(val_loader, resnet18_clustered, criterion, args, verbose=False)\n",
    "        \n",
    "print('large top-1 {:.2f}, small top-1 {:.2f}, top-1 drop {:.2f}, large top-5 {:.2f}, small top-5 {:.2f}, top-5 drop {:.2f}'.format(val_top1,test_top1_acc,val_top1-test_top1_acc,val_top5,test_top5_acc,val_top5-test_top5_acc))       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  + Number of FLOPs: 3.64G\n",
      "  + Number of FLOPs: 1.97G\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1970388704.0"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from src.compute_flops import print_model_param_nums,print_model_param_flops\n",
    "\n",
    "print_model_param_flops(resnet18.cpu(),input_res=224)\n",
    "print_model_param_flops(resnet18_clustered.module.cpu(),input_res=224)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Try pruning alongside training the model (CUP-SS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Summary\n",
    "\n",
    "1. K  = 0.03, b = 0.3 : small top-1 67.38,  small top-5 87.86\n",
    "  \n",
    "  + Number of FLOPs: 3.64G\n",
    "  + Number of FLOPs: 2.09G\n",
    "\n",
    "2. K  = 0.03, b = 0.4 : small top-1 66.86,  small top-5 87.37\n",
    "  \n",
    "  + Number of FLOPs: 3.64G\n",
    "  + Number of FLOPs: 1.92G\n",
    "    \n",
    "3. K  = 0.03, b = 0.5 : small top-1 67.24,  small top-5 87.59\n",
    "  \n",
    "  + Number of FLOPs: 3.64G\n",
    "  + Number of FLOPs: 1.99G"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### K  = 0.03, b = 0.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.compute_flops import print_model_param_nums,print_model_param_flops\n",
    "\n",
    "set_random_seed(args.seed)\n",
    "\n",
    "resnet18 = torchvision.models.resnet18(pretrained=False)\n",
    "resnet18_clustered = torch.nn.DataParallel(resnet18)\n",
    "resnet18_clustered.cuda(args.gpu)\n",
    "\n",
    "cluster_args = {    \n",
    "    'cluster_layers' : {4:0,9:0,14:0,21:0,26:0,33:0,38:0,45:0},\n",
    "    'conv_feature_size' : 1,\n",
    "    'features' : 'both',\n",
    "    'channel_reduction' : 'fro',\n",
    "    'use_bias' : False,\n",
    "    'reshape_exists' : False,\n",
    "    'linkage_method' : 'ward',\n",
    "    'distance_metric' : 'euclidean',\n",
    "    'cluster_criterion' : 'hierarchical',\n",
    "    'distance_threshold' : 0.825,\n",
    "    'merge_criterion' : 'max_l2_norm',\n",
    "    'verbose' : False\n",
    "}\n",
    "\n",
    "path = args.checkpoint_path[:-4] + '_lr_point1_cupSS_K_point03_b_point3.pth' \n",
    "\n",
    "T_values = {}\n",
    "slope,b = 0.03,0.3\n",
    "\n",
    "for epoch in range(args.epochs+1):\n",
    "    T_values[epoch] = slope * (epoch) + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/nethome/rduggal7/.conda/envs/kdd_final/lib/python3.6/site-packages/torch/serialization.py:425: SourceChangeWarning: source code of class 'torch.nn.parallel.data_parallel.DataParallel' has changed. you can retrieve the original source code by accessing the object's source attribute or set `torch.nn.Module.dump_patches = True` and use the patch tool to revert the changes.\n",
      "  warnings.warn(msg, SourceChangeWarning)\n",
      "/nethome/rduggal7/.conda/envs/kdd_final/lib/python3.6/site-packages/torch/serialization.py:425: SourceChangeWarning: source code of class 'torch.nn.modules.conv.Conv2d' has changed. you can retrieve the original source code by accessing the object's source attribute or set `torch.nn.Module.dump_patches = True` and use the patch tool to revert the changes.\n",
      "  warnings.warn(msg, SourceChangeWarning)\n",
      "/nethome/rduggal7/.conda/envs/kdd_final/lib/python3.6/site-packages/torch/serialization.py:425: SourceChangeWarning: source code of class 'torch.nn.modules.batchnorm.BatchNorm2d' has changed. you can retrieve the original source code by accessing the object's source attribute or set `torch.nn.Module.dump_patches = True` and use the patch tool to revert the changes.\n",
      "  warnings.warn(msg, SourceChangeWarning)\n",
      "/nethome/rduggal7/.conda/envs/kdd_final/lib/python3.6/site-packages/torch/serialization.py:425: SourceChangeWarning: source code of class 'torch.nn.modules.linear.Linear' has changed. you can retrieve the original source code by accessing the object's source attribute or set `torch.nn.Module.dump_patches = True` and use the patch tool to revert the changes.\n",
      "  warnings.warn(msg, SourceChangeWarning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " * Acc@1 67.384 Acc@5 87.864\n",
      "small top-1 67.38,  small top-5 87.86\n"
     ]
    }
   ],
   "source": [
    "args.epochs = 90\n",
    "args.lr = 0.1\n",
    "best_val_acc = 0\n",
    "flag = True\n",
    "\n",
    "\n",
    "if not os.path.isfile(path):\n",
    "    for epoch in range(1, args.epochs+1):          \n",
    "        if epoch in T_values.keys() and flag:\n",
    "            print('changing T value to {}'.format(T_values[epoch]))\n",
    "            cluster_args['distance_threshold'] = T_values[epoch]        \n",
    "            model_modifier = cluster_model(resnet18_clustered.module,cluster_args)\n",
    "            resnet18_clustered = model_modifier.cluster_model()            \n",
    "            resnet18_clustered = torch.nn.DataParallel(resnet18_clustered)\n",
    "            resnet18_clustered.cuda(args.gpu)\n",
    "            \n",
    "            optimizer = torch.optim.SGD(resnet18_clustered.parameters(), args.lr,\n",
    "                                        momentum=args.momentum,\n",
    "                                        weight_decay=args.weight_decay)\n",
    "            T = T_values[epoch]            \n",
    "            flops = print_model_param_flops(copy.deepcopy(resnet18_clustered.module).cpu(),input_res=224)\n",
    "            if (flops/1e9 <= 2.12):\n",
    "                print('stop filter pruning')\n",
    "                flag = False\n",
    "                \n",
    "                \n",
    "        adjust_learning_rate_pytorch_retrain(optimizer, epoch, args)        \n",
    "\n",
    "        # train for one epoch\n",
    "        train_loss,train_top1,train_top5 = train(train_loader, resnet18_clustered, criterion, optimizer, epoch, args)\n",
    "\n",
    "#         evaluate on validation set\n",
    "        val_loss,val_top1,val_top5 = validate(val_loader, resnet18_clustered, criterion, args)\n",
    "        \n",
    "        if val_top1 > best_val_acc:  \n",
    "            torch.save(resnet18_clustered, path, pickle_protocol=4)            \n",
    "            best_val_acc = val_top1    \n",
    "\n",
    "        writer.add_scalars('resnet18_imagenet_lr_point1_k_point_03_b_point3/loss',{'train_loss': train_loss,\n",
    "                                        'val_loss' : val_loss}, epoch)\n",
    "        writer.add_scalars('resnet18_imagenet_lr_point1_k_point_03_b_point3/accuracy',{'train_top1': train_top1,\n",
    "                                                  'val_top1': val_top1,\n",
    "                                                  'train_top5': train_top5,\n",
    "                                                  'val_top5': val_top5}, epoch) \n",
    "       \n",
    "        writer.add_scalars('resnet18_imagenet_lr_point1_k_point_03_b_point3/flops',{'flops': flops}, epoch) \n",
    "        writer.add_scalars('resnet18_imagenet_lr_point1_k_point_03_b_point3/T',{'T': T}, epoch) \n",
    "else:\n",
    "    resnet18_clustered = torch.load(path).module\n",
    "    test_loss,test_top1_acc,test_top5_acc = validate(val_loader, resnet18_clustered, criterion, args, verbose=False)\n",
    "        \n",
    "print('small top-1 {:.2f},  small top-5 {:.2f}'.format(test_top1_acc,test_top5_acc))        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  + Number of FLOPs: 3.64G\n",
      "  + Number of FLOPs: 2.09G\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2088542600.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from src.compute_flops import print_model_param_nums,print_model_param_flops\n",
    "\n",
    "print_model_param_flops(resnet18.cpu(),input_res=224)\n",
    "print_model_param_flops(resnet18_clustered.cpu(),input_res=224)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6 (kdd_final)",
   "language": "python",
   "name": "kdd_final"
  },
  "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}