vehicles_on_snow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os, warnings\n",
    "from argparse import Namespace\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import torch as ch\n",
    "from torchvision import transforms\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "\n",
    "from helpers import classifier_helpers\n",
    "import helpers.data_helpers as dh\n",
    "import helpers.context_helpers as coh\n",
    "import helpers.rewrite_helpers as rh\n",
    "import helpers.vis_helpers as vh\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_NAME = 'ImageNet'\n",
    "LAYERNUM = 12\n",
    "REWRITE_MODE = 'editing'\n",
    "ARCH = 'vgg16'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = classifier_helpers.get_default_paths(DATASET_NAME, arch=ARCH)\n",
    "DATASET_PATH, MODEL_PATH, MODEL_CLASS, ARCH, CD = ret\n",
    "ret = classifier_helpers.load_classifier(MODEL_PATH, MODEL_CLASS, ARCH,\n",
    "                            DATASET_NAME, LAYERNUM) \n",
    "model, context_model, target_model = ret[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load base dataset and vehicles-on-snow data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dataset, train_loader, val_loader = dh.get_dataset(DATASET_NAME, DATASET_PATH,\n",
    "                                                        batch_size=32, workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, test_data = dh.get_vehicles_on_snow_data(DATASET_NAME, CD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Train exemplars\")\n",
    "vh.show_image_row([train_data['imgs'], train_data['masks'], train_data['modified_imgs']], \n",
    "                  ['Original', 'Mask', 'Modified'], fontsize=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Flickr-sourced test set\")\n",
    "for c, x in test_data.items():\n",
    "    vh.show_image_row([x[:5]], title=f'{CD[c]} ({c})')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate model performance on test set pre-rewriting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "print(\"Original accuracy on test vehicles-on-snow data\")\n",
    "\n",
    "RESULTS = {k: {'preds': {}, 'acc': {}} for k in ['pre', 'post']}\n",
    "for c, x in test_data.items():\n",
    "    with ch.no_grad():\n",
    "        pred = model(x.cuda()).argmax(axis=1)\n",
    "    correct = [p for p in pred if p == c]\n",
    "    acc = 100 * len(correct) / len(x)\n",
    "    print(f'Class: {c}/{CD[c]} | Accuracy: {acc:.2f}',) \n",
    "    RESULTS['pre']['acc'][c] = acc\n",
    "    RESULTS['pre']['preds'][c] = pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform re-write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_args = {'ntrain': 1, # Number of exemplars\n",
    "            'arch': ARCH, # Network architecture\n",
    "            'mode_rewrite': REWRITE_MODE, # Rewriting method ['editing', 'finetune_local', 'finetune_global']\n",
    "            'layernum': LAYERNUM, # Layer to modify\n",
    "            'nsteps': 20000 if REWRITE_MODE == 'editing' else 400, # Number of rewriting steps  \n",
    "            'lr': 1e-4, # Learning rate\n",
    "            'restrict_rank': True, # Whether or not to perform low-rank update\n",
    "            'nsteps_proj': 10, # Frequency of weight projection\n",
    "            'rank': 1, # Rank of subspace to project weights\n",
    "            'use_mask': True # Whether or not to use mask\n",
    "             }\n",
    "train_args = Namespace(**train_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context_model = rh.edit_classifier(train_args, \n",
    "                                   train_data, \n",
    "                                   context_model, \n",
    "                                   target_model=target_model, \n",
    "                                   val_loader=val_loader,\n",
    "                                   caching_dir=f\"./cache/covariances/{DATASET_NAME}_{ARCH}_layer{LAYERNUM}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate model performance on test set post-rewriting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Change in accuracy on test vehicles-on-snow data \\n\")\n",
    "\n",
    "for c, x in test_data.items():\n",
    "    with ch.no_grad():\n",
    "        pred = model(x.cuda()).argmax(axis=1)\n",
    "    correct = [p for p in pred if p == c]\n",
    "    acc = 100 * len(correct) / len(x)\n",
    "    print(f'Class: {c}/{CD[c]} \\n Accuracy change: {RESULTS[\"pre\"][\"acc\"][c]:.2f} -> {acc:.2f} \\n',) \n",
    "    RESULTS['post']['acc'][c] = acc\n",
    "    RESULTS['post']['preds'][c] = pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}