notebooks/playground.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MUFIA attack\n",
    "\n",
    "### Notebook for generating images from the main paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from torchvision.models import *\n",
    "from torchvision.utils import *\n",
    "import torch\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import ipywidgets as widgets\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "# fix torch seed\n",
    "torch.manual_seed(42)\n",
    "# fix cuda seed\n",
    "torch.cuda.manual_seed(42)\n",
    "# fix cudnn seed for reproducibility\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "\n",
    "# Defaults\n",
    "use_cuda = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Parameters\n",
    "\n",
    "Set the dataset and model architecture that will be attacked. The user can also choose various threat models like standard trained model, prime augmented model, etc.\n",
    "The attack parameters can also be set here.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "param = {}\n",
    "param['batch_size'] = 32\n",
    "param['threat_model'] = 'std' # std, prime, augmix, cc_sota\n",
    "param['dataset'] = 'cifar10' # cifar10, cifar100, imagenet\n",
    "\n",
    "param['save_mat'] = False # if true, save the filter\n",
    "param['parallel'] = False\n",
    "\n",
    "if param['dataset'] == 'cifar100':\n",
    "    param['model_name'] = 'resnet56' \n",
    "else:\n",
    "    param['model_name'] = 'resnet50'\n",
    "\n",
    "if param['dataset'] == 'imagenet':\n",
    "    param['block_size'] = 56\n",
    "else:\n",
    "    param['block_size'] = 32\n",
    "\n",
    "param['lambda_reg'] = 20 \n",
    "param['atk_type'] = 'mufia'\n",
    "param['n_epochs'] = 100\n",
    "param['print_every'] = 10\n",
    "param['lr'] = 0.1\n",
    "param['verbose'] = True\n",
    "param['kappa'] = 0.99\n",
    "param[\"sim_loss\"] = 'cosine'\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define dataloaders\n",
    "import sys\n",
    "sys.path.insert(0, \"..\")\n",
    "sys.path.append(\"../\")\n",
    "from misc import *\n",
    "from data import *\n",
    "from models import *\n",
    "\n",
    "data_loading = DataLoading(params=param)\n",
    "_, _, testset, testloader = data_loading.get_data()\n",
    "\n",
    "param[\"dataloader\"] = testloader\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Model Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_loading = ModelLoader(params=param, device=device)\n",
    "net = model_loading.get_model()\n",
    "\n",
    "net = net.to(device)\n",
    "net.eval()\n",
    "# dont print \n",
    "net.verbose = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adv_acc = {}\n",
    "clean_acc = {}\n",
    "y_quantize = {}\n",
    "adv_x = {}\n",
    "adv_acc['std'] = 0\n",
    "y_quantize['std'] = []\n",
    "adv_x['std'] = []\n",
    "corrupt_x = []\n",
    "clean_acc['std'] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from attacks.attacks import *\n",
    "from attacks.attack_utils import *"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running MUFIA on one batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_solver = FilterAttack(net, param, device)\n",
    "for i, (x, y) in enumerate(tqdm(testloader)):\n",
    "    x, y = x.to(device), y.to(device)\n",
    "    corrupt_x.append(x)\n",
    "    adv_x_basic, y_quantize_basic = new_solver(x, y)\n",
    "    adv_x['std'].append(adv_x_basic)\n",
    "\n",
    "    # get adv acc\n",
    "    new_outputs = net(adv_x_basic)\n",
    "    new_pred = (\n",
    "        new_outputs[0] if (type(new_outputs) is tuple) else new_outputs\n",
    "    ).argmax(dim=1, keepdim=True)\n",
    "    adv_acc_curr = new_pred.eq(y.view_as(new_pred)).sum().item()\n",
    "    adv_acc['std'] += adv_acc_curr\n",
    "\n",
    "    # get clean acc\n",
    "    old_outputs = net(x)\n",
    "    old_pred = (\n",
    "        old_outputs[0] if (type(old_outputs) is tuple) else old_outputs\n",
    "    ).argmax(dim=1, keepdim=True)\n",
    "    clean_acc_curr = old_pred.eq(y.view_as(old_pred)).sum().item()\n",
    "    clean_acc['std'] += clean_acc_curr\n",
    "\n",
    "    y_quantize['std'].append(y_quantize_basic)\n",
    "    \n",
    "    # comment this out if you want to run for all the images\n",
    "    if i == 0:\n",
    "        break\n",
    "\n",
    "adv_acc['std'] = 100.0 * adv_acc['std'] / ((i + 1) * param['batch_size'])\n",
    "clean_acc['std'] = 100.0 * clean_acc['std'] / ((i + 1) * param['batch_size'])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"After MUFIA accuracy: \", adv_acc['std'])\n",
    "print(\"Clean accuracy: \", clean_acc['std'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@widgets.interact(idx=(0, param['batch_size']-1))\n",
    "def f(idx=0):\n",
    "    # data, targets = x[idx:idx+1], y[idx:idx+1]\\\n",
    "    data, targets = testset[idx]\n",
    "    data = data.to(device).unsqueeze(0)\n",
    "\n",
    "\n",
    "    y_quantize_new = y_quantize['std'][0][idx:idx+1]\n",
    "    # print(y_quantize_new.shape)\n",
    "    \n",
    "    adv_data_new = adv_x['std'][0][idx:idx+1]\n",
    "    \n",
    "    y_quantize_new = y_quantize_new[0, 0, 0, :, :]\n",
    "\n",
    "    y_quantize_new = torch.abs(y_quantize_new)\n",
    "    y_quantize_new = torch.tanh(y_quantize_new-1)\n",
    "\n",
    "    img_diff = adv_data_new - data\n",
    "    img_diff = torch.abs(img_diff)\n",
    "    img_diff = img_diff * 10 # scale it up\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 5)\n",
    "    fig.set_size_inches(10, 10)\n",
    "\n",
    "    # remove the border\n",
    "    for a in ax:\n",
    "        a.axis('off')\n",
    "    ax[0].imshow(data.cpu().squeeze().numpy().transpose(1, 2, 0))\n",
    "    ax[1].imshow(data.cpu().squeeze().numpy().transpose(1, 2, 0))\n",
    "    ax[2].imshow(adv_data_new.cpu().squeeze().numpy().transpose(1, 2, 0))\n",
    "    ax[3].imshow(img_diff.cpu().squeeze().numpy().transpose(1, 2, 0))\n",
    "    import matplotlib as mpl\n",
    "    cmap = plt.get_cmap('tab20c')\n",
    "    norm = mpl.colors.Normalize(vmin=-1, vmax=1)\n",
    "    ax[4].imshow(y_quantize_new.cpu().numpy(), cmap=cmap)\n",
    "    pos = ax[4].get_position()\n",
    "    \n",
    "    cbar_ax = fig.add_axes([pos.x0 + pos.width + 0.01, pos.y0, 0.01, pos.height])\n",
    "    fig.colorbar(ax[4].imshow(y_quantize_new.cpu().numpy(),\n",
    "                 cmap=cmap, norm=norm), cax=cbar_ax)\n",
    "    ax[0].set_title(\"True Label: \" + str(targets), y=1.05)\n",
    "    ax[1].set_title(\"Pred Label: \" + str(net(data).argmax().item()), y=1.05)\n",
    "    ax[2].set_title(\"Adv Label: \" + str(net(adv_data_new).argmax().item()), y=1.05)\n",
    "    ax[3].set_title(\"Difference\", y=1.05)\n",
    "    ax[4].set_title(\"Filter Bank\", y=1.05)\n",
    "\n",
    "    # remove axis ticks \n",
    "    for i in range(5):\n",
    "        ax[i].set_xticks([])\n",
    "        ax[i].set_yticks([])\n",
    "    fig.subplots_adjust(wspace=0.1)\n",
    "    plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('machiraj_attks')",
   "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.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c0fc4fd3b2f4f47fd2da61ab20e7060105922dd788f35e5aa10c938a4ba7a375"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}