main.ipynb

{
    "cells":
    [
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "x8ktP4NSN_vZ"
            },
            "source":
            [
                "# Setup"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "SM3DdrGBO-3d"
            },
            "source":
            [
                "## Mount Filesystem"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "zgNChE1K_-cP"
            },
            "outputs":
            [],
            "source":
            [
                "# Set up colab filesystem\n",
                "import os\n",
                "import sys\n",
                "os.environ['NOTEBOOK_MODE'] = '1'\n",
                "from google.colab import drive\n",
                "\n",
                "os.chdir('/content')\n",
                "drive.mount('drive')\n",
                "WORKDIR = 'drive/MyDrive/Colab Notebooks/ECE661/Project' # Change to location in drive\n",
                "sys.path.append(WORKDIR)\n",
                "os.chdir(WORKDIR)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "HQqxyoBcPC-n"
            },
            "source":
            [
                "## Dependencies"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "5HqpTv5COGQk"
            },
            "outputs":
            [],
            "source":
            [
                "# Install python dependencies\n",
                "!pip install -r requirements.txt"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "67SNSTMb_nVx"
            },
            "outputs":
            [],
            "source":
            [
                "# Import dependencies\n",
                "import torch as ch\n",
                "from torch.nn.functional import interpolate\n",
                "from torch.distributions.multivariate_normal import MultivariateNormal\n",
                "import numpy as np\n",
                "import seaborn as sns\n",
                "from scipy import stats\n",
                "from skimage.transform import resize\n",
                "from tqdm import tqdm, tqdm_notebook\n",
                "from PIL import Image\n",
                "import torchvision.transforms as transforms\n",
                "import matplotlib.pyplot as plt\n",
                "from robustness import model_utils, datasets\n",
                "from robustness.tools.vis_tools import show_image_row, show_image_column\n",
                "from robustness.tools.label_maps import CLASS_DICT\n",
                "from user_constants import DATA_PATH_DICT\n",
                "%matplotlib inline"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "3TDQJbd5OPG8"
            },
            "source":
            [
                "## Constants"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "sGRhGkPa_nV1"
            },
            "outputs":
            [],
            "source":
            [
                "# Constants\n",
                "DATA = 'ImageNet' # Choices: ['CIFAR', 'ImageNet'] Use ImageNet for ImageNette\n",
                "IMAGENETTE_CLASSES = [0, 217, 482, 491, 497, 566, 569, 571, 574, 701] # Map index to target class\n",
                "IMAGENETTE_CLASS_NAMES = ['tench', 'english springer', 'casette player', \n",
                "                          'chain saw', 'church', 'french horn', \n",
                "                          'garbage truck', 'gas pump', 'golf ball', 'parachute']\n",
                "BATCH_SIZE = 5\n",
                "NUM_WORKERS = 2\n",
                "NUM_CLASSES_VIS = 10\n",
                "DEVICE = 'cuda'\n",
                "USE_PYTORCH_PRETRAINED = False\n",
                "\n",
                "DATA_SHAPE = 32 if DATA == 'CIFAR' else 224 # Image size (fixed for dataset)\n",
                "REPRESENTATION_SIZE = 2048 # Size of representation vector (fixed for model)\n",
                "CLASSES = CLASS_DICT[DATA] # Class names for dataset\n",
                "NUM_CLASSES = len(CLASSES) - 1 \n",
                "NUM_CLASSES_VIS = min(NUM_CLASSES_VIS, NUM_CLASSES)\n",
                "GRAIN = 4 if DATA != 'CIFAR' else 1\n",
                "CLASS_IDS = IMAGENETTE_CLASSES if DATA == 'ImageNet' else range(NUM_CLASSES_VIS)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "j3UL0pl2OYmx"
            },
            "source":
            [
                "## Load Dataset and Pretrained Model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "IDDd8ptd_nV2"
            },
            "outputs":
            [],
            "source":
            [
                "# Load dataset\n",
                "dataset_function = getattr(datasets, DATA)\n",
                "dataset = dataset_function(DATA_PATH_DICT[DATA])\n",
                "train_loader, test_loader = dataset.make_loaders(workers=NUM_WORKERS, \n",
                "                                      batch_size=BATCH_SIZE, \n",
                "                                      data_aug=False)\n",
                "data_iterator = enumerate(test_loader)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "TDiL_auR_nV3"
            },
            "outputs":
            [],
            "source":
            [
                "# Load model\n",
                "if USE_PYTORCH_PRETRAINED:\n",
                "    model_kwargs = {\n",
                "        'arch': 'resnet50',\n",
                "        'dataset': dataset,\n",
                "        'pytorch_pretrained': True,\n",
                "    }\n",
                "else:\n",
                "    model_kwargs = {\n",
                "        'arch': 'resnet50',\n",
                "        'dataset': dataset,\n",
                "        'resume_path': f'./models/{DATA}.pt',\n",
                "    }\n",
                "\n",
                "model, _ = model_utils.make_and_restore_model(**model_kwargs)\n",
                "model.eval()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "3pjnJdu-_nV4"
            },
            "outputs":
            [],
            "source":
            [
                "def downsample(x, step=GRAIN):\n",
                "    down = ch.zeros([len(x), 3, DATA_SHAPE//step, DATA_SHAPE//step])\n",
                "\n",
                "    for i in range(0, DATA_SHAPE, step):\n",
                "        for j in range(0, DATA_SHAPE, step):\n",
                "            v = x[:, :, i:i+step, j:j+step].mean(dim=2, keepdim=True).mean(dim=3, keepdim=True)\n",
                "            ii, jj = i // step, j // step\n",
                "            down[:, :, ii:ii+1, jj:jj+1] = v\n",
                "    return down\n",
                "\n",
                "def upsample(x, step=GRAIN):\n",
                "    up = ch.zeros([len(x), 3, DATA_SHAPE, DATA_SHAPE])\n",
                "\n",
                "    for i in range(0, DATA_SHAPE, step):\n",
                "        for j in range(0, DATA_SHAPE, step):\n",
                "            ii, jj = i // step, j // step\n",
                "            up[:, :, i:i+step, j:j+step] = x[:, :, ii:ii+1, jj:jj+1]\n",
                "    return up"
            ]
        },
        {
            "cell_type": "code",
            "source":
            [
                "def index_to_class(targ):\n",
                "    # Convert index as read from dataloader to actual target class\n",
                "    targ_list = []\n",
                "    for i in targ:\n",
                "        targ_list.append(CLASS_IDS[i])\n",
                "    return ch.tensor(targ_list)"
            ],
            "metadata":
            {
                "id": "3fHt2OuqIH0Q"
            },
            "execution_count": null,
            "outputs":
            []
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "EhnYYSwAOxns"
            },
            "source":
            [
                "# Image Generation"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "605VKI8cT15s"
            },
            "source":
            [
                "## Calculate Random Seeds"
            ]
        },
        {
            "cell_type": "code",
            "source":
            [
                "# Get seed distribution (can be memory intensive to do all ImageNet classes at once)\n",
                "def get_seed_distributions(imc):\n",
                "    down_flat = downsample(imc).view(len(imc), -1)\n",
                "    mean = down_flat.mean(dim=0)\n",
                "    down_flat = down_flat - mean.unsqueeze(dim=0)\n",
                "    cov = down_flat.t() @ down_flat / len(imc)\n",
                "    return MultivariateNormal(mean, covariance_matrix=cov+1e-4*ch.eye(3 * DATA_SHAPE//GRAIN * DATA_SHAPE//GRAIN))"
            ],
            "metadata":
            {
                "id": "3gb5grkGeGuc"
            },
            "execution_count": null,
            "outputs":
            []
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "oOnrpFMR_nV4"
            },
            "outputs":
            [],
            "source":
            [
                "# Load images\n",
                "conditionals = {}\n",
                "for i in CLASS_IDS:\n",
                "    im_test = []\n",
                "    for _, (im, targ) in enumerate(tqdm(test_loader)):\n",
                "        targets = index_to_class(targ)\n",
                "        for j, t in enumerate(targets):\n",
                "            if t == i:\n",
                "                im_test.append(im[j])\n",
                "    im_test = ch.stack(im_test)\n",
                "    print(f'Loaded {im_test.shape[0]} images for Class {i}')\n",
                "    conditionals[i] = get_seed_distributions(im_test)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "WtrCToDd_nV5",
                "collapsed": true
            },
            "outputs":
            [],
            "source":
            [
                "# Visualize seeds\n",
                "img_seed = ch.stack([conditionals[i].sample().view(3, DATA_SHAPE//GRAIN, DATA_SHAPE//GRAIN) \n",
                "                     for i in CLASS_IDS])\n",
                "img_seed = ch.clamp(img_seed, min=0, max=1)\n",
                "show_image_row([img_seed.cpu()], \n",
                "               tlist=[[CLASSES[int(t)].split(',')[0] for t in CLASS_IDS]], \n",
                "               filename=f'./results/{DATA}/seeds.png')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "8jXffPkGUGJN"
            },
            "source":
            [
                "## Generate Images"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "ZeT33ZOr_nV5"
            },
            "outputs":
            [],
            "source":
            [
                "def generation_loss(mod, inp, targ):\n",
                "    op = mod(inp)\n",
                "    loss = ch.nn.CrossEntropyLoss(reduction='none')(op, targ)\n",
                "    return loss, None\n",
                "\n",
                "kwargs = {\n",
                "        'custom_loss': generation_loss,\n",
                "        'constraint':'2',\n",
                "        'eps': 40,\n",
                "        'step_size': 1,\n",
                "        'iterations': 60,\n",
                "        'targeted': True,\n",
                "        'random_start': False,\n",
                "}  \n",
                "    \n",
                "if DATA == 'CIFAR':\n",
                "    kwargs['eps'] = 30\n",
                "    kwargs['step_size'] = 0.5\n",
                "    kwargs['iterations'] = 60"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "qCCPKlUw_nV6",
                "scrolled": false
            },
            "outputs":
            [],
            "source":
            [
                "show_seed = False\n",
                "USE_SEED = True\n",
                "\n",
                "for i in CLASS_IDS:\n",
                "    target_class = i * ch.ones((BATCH_SIZE, ))\n",
                "    if USE_SEED:\n",
                "      im_seed = ch.stack([conditionals[int(t)].sample().view(3, DATA_SHAPE//GRAIN, DATA_SHAPE//GRAIN) \n",
                "                          for t in target_class])\n",
                "      im_seed = upsample(ch.clamp(im_seed, min=0, max=1)).to(DEVICE)\n",
                "      if USE_PYTORCH_PRETRAINED:\n",
                "          save_path = f'./results/{DATA}/generation/non-robust/{i}.png'\n",
                "      else:\n",
                "          save_path = f'./results/{DATA}/generation/seed/{i}.png'\n",
                "    else:\n",
                "      im_seed = ch.rand((BATCH_SIZE, 3, DATA_SHAPE, DATA_SHAPE)).to(DEVICE)\n",
                "      save_path = f'./results/{DATA}/generation/random/{i}.png'\n",
                "    _, im_gen = model(im_seed, target_class.long().to(DEVICE), make_adv=True, **kwargs)\n",
                "    if show_seed:\n",
                "        show_image_row([im_seed.cpu()], [f'Seed ($x_0$)'], fontsize=18)\n",
                "    show_image_row([im_gen.detach().cpu()], \n",
                "                   [CLASSES[int(t)].split(',')[0] for t in target_class], \n",
                "                   fontsize=18,\n",
                "                   filename=save_path)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "gZn1Hz7cUJxD"
            },
            "source":
            [
                "# Inpainting"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "hp79cUWuU8o0"
            },
            "source":
            [
                "## Image Corruption"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "j3Sezt8MURD2"
            },
            "outputs":
            [],
            "source":
            [
                "PATCH_SIZE = 6 if DATA == \"CIFAR\" else 60\n",
                "\n",
                "# Function to get corrupted image\n",
                "def mask_image(x, width=PATCH_SIZE):\n",
                "    loc = np.random.randint(0, x.shape[-1] - PATCH_SIZE, size=(x.shape[0], 2))\n",
                "    mask = ch.zeros_like(x)\n",
                "    for idx in range(x.shape[0]):\n",
                "        i, j = loc[idx, 0], loc[idx, 1]\n",
                "        val = ch.mean(ch.mean(x[idx, :], dim=2, keepdim=True), dim=1, keepdim=True)\n",
                "        # Initialize masked region as mean pixel value over image (per channel)\n",
                "        x[idx, :, i:i+PATCH_SIZE, j:j+PATCH_SIZE] = val.expand_as(x[idx, :, i:i+PATCH_SIZE, j:j+PATCH_SIZE])\n",
                "        mask[idx, :, i:i+PATCH_SIZE, j:j+PATCH_SIZE] = 1\n",
                "    return x, mask"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "kpp3OXUYVD2l"
            },
            "source":
            [
                "## Inpainting"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "XYpouZXrURD2"
            },
            "outputs":
            [],
            "source":
            [
                "# Custom inpainting loss\n",
                "def inpainting_loss_wrapper(im_targ, mask, normalizer, lambdar=10):\n",
                "    def inpainting_loss(mod, inp, targ):\n",
                "        op = mod(normalizer(inp), fake_relu=True)\n",
                "        loss = ch.nn.CrossEntropyLoss(reduction='none')(op, targ)\n",
                "        loss_l2 = ((im_targ - inp) * (1 - mask))**2 \n",
                "        loss_l2 = loss_l2.mean(-1).mean(-1).mean(-1)\n",
                "        loss += lambdar * loss_l2\n",
                "        return loss, None\n",
                "    return inpainting_loss"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "hoDekJ2WURD3"
            },
            "outputs":
            [],
            "source":
            [
                "# PGD parameters\n",
                "kwargs = {\n",
                "        'constraint':'2',\n",
                "        'eps': 21.6,\n",
                "        'step_size': 0.1,\n",
                "        'iterations': 720, \n",
                "        'do_tqdm': True,\n",
                "        'targeted': True,\n",
                "        'should_normalize': False\n",
                "}"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "JzIKk0x6URD3",
                "scrolled": false
            },
            "outputs":
            [],
            "source":
            [
                "NUM_IMGS_PER_CLASS = 1\n",
                "# Inpainting using PGD\n",
                "images, targets = [], []\n",
                "for _, (im, targ) in enumerate(test_loader):\n",
                "    images.append(im)\n",
                "    targets.append(index_to_class(targ))\n",
                "images, targets = ch.cat(images), ch.cat(targets)\n",
                "\n",
                "img_orig, targ_orig = [], []\n",
                "for i in CLASS_IDS:\n",
                "    target_class = i * ch.ones((NUM_IMGS_PER_CLASS, ), dtype=ch.int64)\n",
                "    imgs = images[targets == i][:NUM_IMGS_PER_CLASS]\n",
                "    img_orig.append(imgs)\n",
                "    targ_orig.append(target_class)\n",
                "\n",
                "img_orig, targ_orig = ch.cat(img_orig), ch.cat(targ_orig)\n",
                "\n",
                "img_corrupt, mask = mask_image(img_orig.clone())\n",
                "img_corrupt = img_corrupt.to(DEVICE)\n",
                "\n",
                "kwargs['custom_loss'] = inpainting_loss_wrapper(img_corrupt.to(DEVICE), \n",
                "                                                mask.to(DEVICE), \n",
                "                                                model.normalizer,\n",
                "                                                lambdar=1000)\n",
                "\n",
                "_, img_inpaint = model(img_corrupt, targ_orig.clone().to(DEVICE), make_adv=True, **kwargs)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "P00HD_DxURD4"
            },
            "outputs":
            [],
            "source":
            [
                "# Visualize inpainted images\n",
                "\n",
                "IMAGES_TO_SHOW = range(NUM_IMGS_PER_CLASS * len(CLASS_IDS))\n",
                "# IMAGES_TO_SHOW = [0, 1, 9]\n",
                "show_corrupt = img_corrupt[IMAGES_TO_SHOW, :, :, :]\n",
                "show_orig = img_orig[IMAGES_TO_SHOW, :, :, :]\n",
                "show_inpaint = img_inpaint[IMAGES_TO_SHOW, :, :, :]\n",
                "show_targ = targ_orig[IMAGES_TO_SHOW]\n",
                "\n",
                "show_image_row([show_corrupt.cpu(), show_orig.cpu(), show_inpaint.detach().cpu()], \n",
                "               ylist=[\"Corrupted\", \"Original\", \"Inpainted\"],\n",
                "               tlist=[[CLASSES[int(t)].split(',')[0] for t in show_targ], [None] * 10, [None] * 10],\n",
                "               fontsize=16,\n",
                "               filename=f'./results/{DATA}/inpainting.png')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "0s-d_GZTaDJh"
            },
            "source":
            [
                "# Super-Resolution"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "2wJQipSXhydk"
            },
            "source":
            [
                "## Generate Downsampled Images"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "w7UdPAQCarDl"
            },
            "outputs":
            [],
            "source":
            [
                "# Get images and classes\n",
                "DOWNSAMPLED_SIZE = (16, 16) if DATA == 'CIFAR' else (64, 64)\n",
                "ORIGINAL_SIZE = (32, 32) if DATA == 'CIFAR' else (224, 224)\n",
                "STEP = 2 if DATA == 'CIFAR' else 7\n",
                "NUM_IMGS_PER_CLASS = 1\n",
                "\n",
                "images, targets = [], []\n",
                "for _, (im, targ) in enumerate(test_loader):\n",
                "    images.append(im)\n",
                "    targets.append(index_to_class(targ))\n",
                "images, targets = ch.cat(images), ch.cat(targets)\n",
                "\n",
                "img, targ = [], []\n",
                "for i in CLASS_IDS:\n",
                "    target_class = i * ch.ones((NUM_IMGS_PER_CLASS, ), dtype=ch.int64)\n",
                "    imgs = images[targets == i][:NUM_IMGS_PER_CLASS]\n",
                "    img.append(imgs)\n",
                "    targ.append(target_class)\n",
                "\n",
                "img, targ = ch.cat(img).to(DEVICE), ch.cat(targ).to(DEVICE)\n",
                "\n",
                "img_downsampled = interpolate(img, size=DOWNSAMPLED_SIZE, mode='nearest').to(DEVICE)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "VvfCeao6lIJS"
            },
            "source":
            [
                "## Bicubic Upsampling"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "qpEPB-w7arDp"
            },
            "outputs":
            [],
            "source":
            [
                "img_bicubic = interpolate(img_downsampled, size=ORIGINAL_SIZE, mode='bicubic').clip(0, 1)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata":
            {
                "id": "qhGY-iTLlLit"
            },
            "source":
            [
                "## Super-Resolution Upsampling"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "OLQpjrV7arDm"
            },
            "outputs":
            [],
            "source":
            [
                "# PGD Parameters\n",
                "kwargs = {\n",
                "        'constraint':'2',\n",
                "        'eps': 15,\n",
                "        'step_size': 1,\n",
                "        'iterations': 50,\n",
                "        'do_tqdm':True,\n",
                "        'targeted': True,\n",
                "        'should_normalize': True\n",
                "}"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "wgluGuf1arDo"
            },
            "outputs":
            [],
            "source":
            [
                "# Superresolution using PGD\n",
                "#img_up = upsample(img_downsampled, step=STEP).to(DEVICE)\n",
                "img_up = interpolate(img_downsampled, size=ORIGINAL_SIZE, mode='nearest').clip(0, 1)\n",
                "_, img_sr = model(img_up, targ, fake_relu=False, make_adv=True, **kwargs)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "5Z2D5XuiarDr"
            },
            "outputs":
            [],
            "source":
            [
                "# Visualize superres results\n",
                "\n",
                "IMAGES_TO_SHOW = range(NUM_IMGS_PER_CLASS * len(CLASS_IDS))\n",
                "# IMAGES_TO_SHOW = [0, 4, 5, 8]\n",
                "show_img = img[IMAGES_TO_SHOW, :, :, :]\n",
                "show_downsampled = img_downsampled[IMAGES_TO_SHOW, :, :, :]\n",
                "show_bicubic = img_bicubic[IMAGES_TO_SHOW, :, :, :]\n",
                "show_sr = img_sr[IMAGES_TO_SHOW, :, :, :]\n",
                "show_targ = targ[IMAGES_TO_SHOW]\n",
                "\n",
                "show_image_row([show_img.cpu(), show_downsampled.cpu(), show_bicubic.cpu(), show_sr.detach().cpu()], \n",
                "               [f'Original\\n{ORIGINAL_SIZE}', f'Downsampled\\n{DOWNSAMPLED_SIZE}', \n",
                "                f'Bicubic\\n{ORIGINAL_SIZE}', f'{STEP}x Superres\\n{ORIGINAL_SIZE}'],\n",
                "               tlist=[[CLASSES[int(t)].split(',')[0] for t in show_targ], [None] * 10, [None] * 10, [None] * 10],\n",
                "               fontsize=16,\n",
                "               filename=f'./results/{DATA}/superres.png')"
            ]
        },
        {
            "cell_type": "markdown",
            "source":
            [
                "# Image-to-Image Translation"
            ],
            "metadata":
            {
                "id": "tq_ZJXOSPwsW"
            }
        },
        {
            "cell_type": "markdown",
            "source":
            [
                "## Parameters"
            ],
            "metadata":
            {
                "id": "aUJtQm1LSvW6"
            }
        },
        {
            "cell_type": "code",
            "source":
            [
                "# PGD parameters\n",
                "kwargs = {\n",
                "        #'criterion': ch.nn.CrossEntropyLoss(reduction='none'),\n",
                "        'constraint':'2',\n",
                "        'eps': 80,\n",
                "        'step_size': 0.5,\n",
                "        'iterations': 200,\n",
                "        'do_tqdm': True,\n",
                "        'targeted': True,\n",
                "}"
            ],
            "metadata":
            {
                "id": "4Mx8Y4zJP0r0"
            },
            "execution_count": null,
            "outputs":
            []
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "5NDb_12N6t-g"
            },
            "outputs":
            [],
            "source":
            [
                "# Get random images from dataset\n",
                "NUM_IMGS_PER_CLASS = 1\n",
                "\n",
                "images, targets = [], []\n",
                "for _, (im, targ) in enumerate(test_loader):\n",
                "    images.append(im)\n",
                "    targets.append(index_to_class(targ))\n",
                "images, targets = ch.cat(images), ch.cat(targets)\n",
                "\n",
                "img, targ = [], []\n",
                "for i in CLASS_IDS:\n",
                "    target_class = i * ch.ones((NUM_IMGS_PER_CLASS, ), dtype=ch.int64)\n",
                "    imgs = images[targets == i][:NUM_IMGS_PER_CLASS]\n",
                "    img.append(imgs)\n",
                "    targ.append(target_class)\n",
                "\n",
                "img, targ = ch.cat(img).to(DEVICE), ch.cat(targ).to(DEVICE)\n",
                "\n",
                "diff_targ = index_to_class(ch.randint(0, 10, (10,))).to(DEVICE)\n",
                "# Perform image-to-image translation using PGD"
            ]
        },
        {
            "cell_type": "markdown",
            "source":
            [
                "## Translate Images"
            ],
            "metadata":
            {
                "id": "zFPr2IBkSyVy"
            }
        },
        {
            "cell_type": "code",
            "source":
            [
                "_, img_translated = model(img, diff_targ, make_adv=True, **kwargs)"
            ],
            "metadata":
            {
                "id": "J7NwHe_hL2dS"
            },
            "execution_count": null,
            "outputs":
            []
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "3Q-moWiR6t-g"
            },
            "outputs":
            [],
            "source":
            [
                "IMAGES_TO_SHOW = range(NUM_IMGS_PER_CLASS * len(CLASS_IDS))\n",
                "# IMAGES_TO_SHOW = [2, 6, 8]\n",
                "show_img = img[IMAGES_TO_SHOW, :, :, :]\n",
                "show_translated = img_translated[IMAGES_TO_SHOW, :, :, :]\n",
                "show_targ = targ[IMAGES_TO_SHOW]\n",
                "show_diff_targ = diff_targ[IMAGES_TO_SHOW]\n",
                "\n",
                "show_image_row([show_img.cpu(), show_translated.detach().cpu()], [\"Original\", \"Translated\"], \n",
                "               tlist=[[CLASSES[int(t)].split(',')[0] for t in show_targ],\n",
                "                    [CLASSES[int(t)].split(',')[0] for t in show_diff_targ]],\n",
                "               fontsize=16,\n",
                "               filename=f'./results/{DATA}/translation.png')"
            ]
        },
        {
            "cell_type": "markdown",
            "source":
            [
                "# Sketch-to-Image Translation"
            ],
            "metadata":
            {
                "id": "51JsRLmYT_i5"
            }
        },
        {
            "cell_type": "markdown",
            "source":
            [
                "## Load Images"
            ],
            "metadata":
            {
                "id": "xv_cGxVzaAKK"
            }
        },
        {
            "cell_type": "code",
            "source":
            [
                "SKETCH_FOLDER = './data/sketches/'\n",
                "IMAGE_SIZE = (32, 32) if DATA == 'CIFAR' else (224, 224)\n",
                "\n",
                "images = []\n",
                "transform = transforms.Compose([\n",
                "    transforms.PILToTensor(),\n",
                "    transforms.Resize(256),\n",
                "    transforms.CenterCrop(IMAGE_SIZE[0])\n",
                "])\n",
                "for i in range(len(CLASS_IDS)):\n",
                "    image = Image.open(os.path.join(SKETCH_FOLDER, f'{i}.png'))\n",
                "    images.append(transform(image))\n",
                "sketch = (ch.stack(images) / 255).to(DEVICE)\n",
                "targ = index_to_class(ch.tensor(range(len(CLASS_IDS)))).to(DEVICE)"
            ],
            "metadata":
            {
                "id": "wlLwsyjpZ_AM"
            },
            "execution_count": null,
            "outputs":
            []
        },
        {
            "cell_type": "markdown",
            "source":
            [
                "## Parameters"
            ],
            "metadata":
            {
                "id": "I5G8aq9bYACE"
            }
        },
        {
            "cell_type": "code",
            "source":
            [
                "kwargs = {\n",
                "        'constraint':'2',\n",
                "        'eps': 40,\n",
                "        'step_size': 1,\n",
                "        'iterations': 60, \n",
                "        'targeted': True,\n",
                "        'should_normalize': True\n",
                "}"
            ],
            "metadata":
            {
                "id": "wETJDuiERF7m"
            },
            "execution_count": null,
            "outputs":
            []
        },
        {
            "cell_type": "markdown",
            "source":
            [
                "## Translate Images"
            ],
            "metadata":
            {
                "id": "YjqUc1q6YCVs"
            }
        },
        {
            "cell_type": "code",
            "source":
            [
                "_, img = model(sketch, targ, make_adv=True, **kwargs)"
            ],
            "metadata":
            {
                "id": "aOVWsSgbWXP7"
            },
            "execution_count": null,
            "outputs":
            []
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata":
            {
                "id": "vJ_VwS3r6t-g"
            },
            "outputs":
            [],
            "source":
            [
                "IMAGES_TO_SHOW = range(NUM_IMGS_PER_CLASS * len(CLASS_IDS))\n",
                "# IMAGES_TO_SHOW = [0, 4, 8]\n",
                "show_img = img[IMAGES_TO_SHOW, :, :, :]\n",
                "show_sketch = sketch[IMAGES_TO_SHOW, :, :, :]\n",
                "show_targ = targ[IMAGES_TO_SHOW]\n",
                "\n",
                "show_image_row([show_sketch.cpu(), show_img.detach().cpu()], [\"Sketch\", \"Translated\"],\n",
                "               tlist=[[CLASSES[int(t)].split(',')[0] for t in show_targ], [None] * 10],\n",
                "               fontsize=16,\n",
                "               filename=f'./results/{DATA}/sketch.png')"
            ]
        }
    ],
    "metadata":
    {
        "accelerator": "GPU",
        "colab":
        {
            "collapsed_sections":
            [
                "x8ktP4NSN_vZ",
                "SM3DdrGBO-3d",
                "HQqxyoBcPC-n",
                "3TDQJbd5OPG8",
                "j3UL0pl2OYmx",
                "EhnYYSwAOxns",
                "605VKI8cT15s",
                "8jXffPkGUGJN",
                "gZn1Hz7cUJxD",
                "hp79cUWuU8o0",
                "kpp3OXUYVD2l",
                "0s-d_GZTaDJh",
                "2wJQipSXhydk",
                "VvfCeao6lIJS",
                "qhGY-iTLlLit",
                "tq_ZJXOSPwsW",
                "aUJtQm1LSvW6",
                "zFPr2IBkSyVy",
                "51JsRLmYT_i5",
                "xv_cGxVzaAKK",
                "I5G8aq9bYACE",
                "YjqUc1q6YCVs"
            ],
            "provenance":
            [],
            "toc_visible": true
        },
        "gpuClass": "standard",
        "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": 0
}