combine_models.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.data\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "\n",
    "import re\n",
    "\n",
    "import os\n",
    "import copy\n",
    "\n",
    "from helpers import makedir, find_high_activation_crop\n",
    "import model\n",
    "import push\n",
    "#import train_and_test as tnt\n",
    "import time\n",
    "import save\n",
    "from log import create_logger\n",
    "from preprocess import mean, std, preprocess_input_function, undo_preprocess_input_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class PPNet_ensemble(nn.Module):\n",
    "    \n",
    "    def __init__(self, ppnets):\n",
    "        \n",
    "        super(PPNet_ensemble, self).__init__()\n",
    "        self.ppnets = ppnets # a list of ppnets\n",
    "    \n",
    "    def forward(self, x):\n",
    "        logits, min_distances_0 = self.ppnets[0](x)\n",
    "        min_distances = [min_distances_0]\n",
    "        for i in range(1, len(self.ppnets)):\n",
    "            logits_i, min_distances_i = self.ppnets[i](x)\n",
    "            logits.add_(logits_i)\n",
    "            min_distances.append(min_distances_i)\n",
    "        return logits, min_distances\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load model from ./saved_models/densenet121/003/30_18push0.8043.pth\n",
      "load model from ./saved_models/resnet34/002/10_19push0.7920.pth\n",
      "load model from ./saved_models/vgg19/003/10_18push0.7822.pth\n",
      "test set size: 5794\n"
     ]
    }
   ],
   "source": [
    "##### MODEL AND DATA LOADING\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
    "\n",
    "# load the models\n",
    "# provide paths to saved models you want to combine:\n",
    "# e.g. load_model_paths = ['./saved_models/densenet121/003/30_18push0.8043.pth',\n",
    "#                          './saved_models/resnet34/002/10_19push0.7920.pth',\n",
    "#                          './saved_models/vgg19/003/10_18push0.7822.pth']\n",
    "# MUST NOT BE EMPTY\n",
    "load_model_paths = []\n",
    "\n",
    "ppnets = []\n",
    "epoch_number_strs = []\n",
    "start_epoch_numbers = []\n",
    "\n",
    "for load_model_path in load_model_paths:\n",
    "    load_model_name = load_model_path.split('/')[-1]\n",
    "    epoch_number_str = re.search(r'\\d+', load_model_name).group(0)\n",
    "    epoch_number_strs.append(epoch_number_str)\n",
    "    \n",
    "    start_epoch_number = int(epoch_number_str)\n",
    "    start_epoch_numbers.append(start_epoch_number)\n",
    "\n",
    "    print('load model from ' + load_model_path)\n",
    "    ppnet = torch.load(load_model_path)\n",
    "    ppnet = ppnet.cuda()\n",
    "    ppnets.append(ppnet)\n",
    "\n",
    "ppnet_ensemble = PPNet_ensemble(ppnets)\n",
    "ppnet_ensemble = ppnet_ensemble.cuda()\n",
    "ppnet_ensemble_multi = torch.nn.DataParallel(ppnet_ensemble)\n",
    "\n",
    "img_size = ppnets[0].img_size\n",
    "\n",
    "#ppnet_multi = torch.nn.DataParallel(ppnet)\n",
    "#img_size = ppnet_multi.module.img_size\n",
    "#prototype_shape = ppnet.prototype_shape\n",
    "#max_dist = prototype_shape[1] * prototype_shape[2] * prototype_shape[3]\n",
    "\n",
    "# load the (test) data\n",
    "from settings import test_dir\n",
    "test_batch_size = 100\n",
    "\n",
    "normalize = transforms.Normalize(mean=mean,\n",
    "                                 std=std)\n",
    "\n",
    "test_dataset = datasets.ImageFolder(\n",
    "    test_dir,\n",
    "    transforms.Compose([\n",
    "        transforms.Resize(size=(img_size, img_size)),\n",
    "        transforms.ToTensor(),\n",
    "        normalize,\n",
    "    ]))\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    test_dataset, batch_size=test_batch_size, shuffle=True,\n",
    "    num_workers=4, pin_memory=False)\n",
    "print('test set size: {0}'.format(len(test_loader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PPNet(\n",
      "\tfeatures: densenet121_features,\n",
      "\timg_size: 224,\n",
      "\tprototype_shape: (2000, 128, 1, 1),\n",
      "\tproto_layer_rf_info: [7, 32, 2071, 14.5],\n",
      "\tnum_classes: 200,\n",
      "\tepsilon: 0.0001\n",
      ")\n",
      "PPNet(\n",
      "\tfeatures: resnet34_features,\n",
      "\timg_size: 224,\n",
      "\tprototype_shape: (2000, 256, 1, 1),\n",
      "\tproto_layer_rf_info: [7, 32, 899, 0.5],\n",
      "\tnum_classes: 200,\n",
      "\tepsilon: 0.0001\n",
      ")\n",
      "PPNet(\n",
      "\tfeatures: VGG19, batch_norm=False,\n",
      "\timg_size: 224,\n",
      "\tprototype_shape: (2000, 128, 1, 1),\n",
      "\tproto_layer_rf_info: [7, 32, 268, 16.0],\n",
      "\tnum_classes: 200,\n",
      "\tepsilon: 0.0001\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "for ppnet in ppnet_ensemble_multi.module.ppnets:\n",
    "    print(ppnet)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class_specific = True\n",
    "\n",
    "# only supports last layer adjustment\n",
    "def _train_or_test_ppnet_ensemble(model, dataloader, optimizer=None, class_specific=True, use_l1_mask=True,\n",
    "                                  coefs=None, log=print):\n",
    "    '''\n",
    "    model: the multi-gpu model\n",
    "    dataloader:\n",
    "    optimizer: if None, will be test evaluation\n",
    "    '''\n",
    "    is_train = optimizer is not None\n",
    "    start = time.time()\n",
    "    n_examples = 0\n",
    "    n_correct = 0\n",
    "    n_batches = 0\n",
    "    total_cross_entropy = 0\n",
    "    \n",
    "    for i, (image, label) in enumerate(dataloader):\n",
    "        input = image.cuda()\n",
    "        target = label.cuda()\n",
    "\n",
    "        # torch.enable_grad() has no effect outside of no_grad()\n",
    "        grad_req = torch.enable_grad() if is_train else torch.no_grad()\n",
    "        with grad_req:\n",
    "            # nn.Module has implemented __call__() function\n",
    "            # so no need to call .forward\n",
    "            output, _ = model(input)\n",
    "\n",
    "            # compute loss\n",
    "            cross_entropy = torch.nn.functional.cross_entropy(output, target)\n",
    "            l1 = torch.tensor(0.0)\n",
    "\n",
    "            if class_specific:\n",
    "                if use_l1_mask:\n",
    "                    for ppnet in model.module.ppnets:\n",
    "                        l1_mask = 1 - torch.t(ppnet.prototype_class_identity).cuda()\n",
    "                        l1_ = (ppnet.last_layer.weight * l1_mask).norm(p=1)\n",
    "                        l1.add_(l1_)\n",
    "                else:\n",
    "                    for ppnet in model.module.ppnets:\n",
    "                        l1_ = ppnet.last_layer.weight.norm(p=1)\n",
    "                        l1.add_(l1_)\n",
    "\n",
    "            else:\n",
    "                for ppnet in model.module.ppnets:\n",
    "                    l1_ = ppnet.last_layer.weight.norm(p=1)\n",
    "                    l1.add_(l1_)\n",
    "\n",
    "            # evaluation statistics\n",
    "            _, predicted = torch.max(output.data, 1)\n",
    "            n_examples += target.size(0)\n",
    "            n_correct += (predicted == target).sum().item()\n",
    "\n",
    "            n_batches += 1\n",
    "            total_cross_entropy += cross_entropy.item()\n",
    "            \n",
    "        # compute gradient and do SGD step\n",
    "        if is_train:\n",
    "            if class_specific:\n",
    "                if coefs is not None:\n",
    "                    loss = (coefs['crs_ent'] * cross_entropy\n",
    "                          + coefs['l1'] * l1)\n",
    "                else:\n",
    "                    loss = cross_entropy + 1e-4 * l1\n",
    "            else:\n",
    "                if coefs is not None:\n",
    "                    loss = (coefs['crs_ent'] * cross_entropy\n",
    "                          + coefs['l1'] * l1)\n",
    "                else:\n",
    "                    loss = cross_entropy + 1e-4 * l1\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        del input\n",
    "        del target\n",
    "        del output\n",
    "        del predicted\n",
    "\n",
    "    end = time.time()\n",
    "\n",
    "    log('\\ttime: \\t{0}'.format(end -  start))\n",
    "    log('\\tcross ent: \\t{0}'.format(total_cross_entropy / n_batches))\n",
    "    log('\\taccu: \\t\\t{0}%'.format(n_correct / n_examples * 100))\n",
    "    last_layer_p1_norm = 0\n",
    "    for ppnet in model.module.ppnets:\n",
    "        last_layer_p1_norm += ppnet.last_layer.weight.norm(p=1).item()\n",
    "    log('\\tl1: \\t\\t{0}'.format(last_layer_p1_norm))\n",
    "    #p = model.module.prototype_vectors.view(model.module.num_prototypes, -1).cpu()\n",
    "    #with torch.no_grad():\n",
    "    #    p_avg_pair_dist = torch.mean(list_of_distances(p, p))\n",
    "    #log('\\tp dist pair: \\t{0}'.format(p_avg_pair_dist.item()))\n",
    "\n",
    "    return n_correct / n_examples\n",
    "\n",
    "def train_ensemble(model, dataloader, optimizer, class_specific=True, coefs=None, log=print):\n",
    "    assert(optimizer is not None)\n",
    "    \n",
    "    log('\\ttrain')\n",
    "    model.train()\n",
    "    return _train_or_test_ppnet_ensemble(model=model, dataloader=dataloader, optimizer=optimizer,\n",
    "                                         class_specific=class_specific, coefs=coefs, log=log)\n",
    "\n",
    "\n",
    "def test_ensemble(model, dataloader, class_specific=True, log=print):\n",
    "    log('\\ttest')\n",
    "    model.eval()\n",
    "    return _train_or_test_ppnet_ensemble(model=model, dataloader=dataloader, optimizer=None,\n",
    "                                         class_specific=class_specific, log=log)\n",
    "\n",
    "\n",
    "def ensemble_last_only(model, log=print):\n",
    "    for ppnet in model.module.ppnets:\n",
    "        for p in ppnet.features.parameters():\n",
    "            p.requires_grad = False\n",
    "        for p in ppnet.add_on_layers.parameters():\n",
    "            p.requires_grad = False\n",
    "        ppnet.prototype_vectors.requires_grad = False\n",
    "        for p in ppnet.last_layer.parameters():\n",
    "            p.requires_grad = True\n",
    "    log('\\tensemble last layer')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\ttest\n",
      "\ttime: \t40.21596312522888\n",
      "\tcross ent: \t1.9900891976109867\n",
      "\taccu: \t\t84.76009665170866%\n",
      "\tl1: \t\t6848.8984375\n"
     ]
    }
   ],
   "source": [
    "#check test accuracy\n",
    "accu = test_ensemble(model=ppnet_ensemble_multi, dataloader=test_loader,\n",
    "                     class_specific=class_specific, log=print)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}