NADs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neural Anisotropy Directions\n",
    "\n",
    "**Authors**: Guillermo Ortiz-Jimenez, Apostolos Modas, Seyed-Mohsen Moosavi-Dezfooli and Pascal Frossard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Requirements"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For executing the code, please make sure that you meet the following requirements:\n",
    "\n",
    "* python (Successfully tested on v3.8.3)\n",
    "* [PyTorch](https://pytorch.org/get-started/previous-versions/) (Successfully tested on v1.5.0 with CUDA v10.0.130)\n",
    "* [Torchvision](https://pytorch.org/get-started/previous-versions/) (Successfully tested on v0.6.0 with CUDA v10.0.130)\n",
    "* [NumPy](https://numpy.org/) (Successfully tested on v1.18.1)\n",
    "* [Matplotlib](https://matplotlib.org/) (Successfully tested on v3.1.3)\n",
    "* [Seaborn](http://seaborn.pydata.org/) (Successfully tested on v0.10.1)\n",
    "* [Scikit-Learn](https://scikit-learn.org/stable/) (Successfully tested on v0.22.1)\n",
    "\n",
    "In our experiments, every package was installed through a Conda environment. Assuming CUDA v10.0.130 and Conda v4.8.1 (installed through [Miniconda3](https://docs.conda.io/en/latest/miniconda.html) on CentOS Linux 7), these are the corresponding commands:\n",
    "\n",
    "```conda create -n nads python==3.8.3```  \n",
    "```conda activate nads```  \n",
    "```conda install numpy==1.18.1```  \n",
    "```conda install pytorch=1.5.0 torchvision=0.6.0 cudatoolkit=10.1 -c pytorch```  \n",
    "```conda install matplotlib==3.1.3```  \n",
    "```conda install seaborn==0.10.1```  \n",
    "```conda install scikit-learn==0.22.1```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of contents\n",
    "\n",
    "- [General training setup](#training_setup)\n",
    "- [Linearly separable dataset](#gen_linear_data)\n",
    "- [NAD Computation](#vis_NADs)\n",
    "- [Poisoning CIFAR-10 dataset](#poison)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data as data\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.decomposition import PCA\n",
    "from torch.utils.data import DataLoader\n",
    "from models import TransformLayer\n",
    "from models import LogReg, LeNet, VGG11_bn, ResNet18, DenseNet121\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a name=training_setup>General training setup</a>\n",
    "\n",
    "We first give the implementation of our main training procedure. Specifically, we use a standard SGD optimizer with a linear learning rate schedule to optimize all networks. The main hyperparameters are:\n",
    "- Number of training epochs:`epochs`\n",
    "- Maximum learning rate:`max_lr`\n",
    "- Momentum:`momentum`\n",
    "- Weight decay:`weight_decay`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, trans, trainloader, testloader,epochs, max_lr, momentum, weight_decay):\n",
    "    lr_schedule = lambda t: np.interp([t], [0, epochs], [max_lr, 0])[0]\n",
    "    opt = torch.optim.SGD(model.parameters(), lr=max_lr, momentum=momentum, weight_decay=weight_decay)\n",
    "    loss_fun = nn.CrossEntropyLoss()\n",
    "    \n",
    "    print('Starting training...')\n",
    "    print()\n",
    "    \n",
    "    best_acc = 0\n",
    "    for epoch in range(epochs):\n",
    "        print('Epoch', epoch)\n",
    "        train_loss_sum = 0\n",
    "        train_acc_sum = 0\n",
    "        train_n = 0\n",
    "\n",
    "        model.train()\n",
    "\n",
    "        for batch_idx, (inputs, targets) in enumerate(trainloader):\n",
    "            inputs, targets = inputs.to(DEVICE), targets.to(DEVICE)\n",
    "\n",
    "            lr = lr_schedule(epoch + (batch_idx + 1) / len(trainloader))\n",
    "            opt.param_groups[0].update(lr=lr)\n",
    "\n",
    "            output = model(trans(inputs))\n",
    "            loss = loss_fun(output, targets)\n",
    "\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), 0.5)\n",
    "            opt.step()\n",
    "\n",
    "            train_loss_sum += loss.item() * targets.size(0)\n",
    "            train_acc_sum += (output.max(1)[1] == targets).sum().item()\n",
    "            train_n += targets.size(0)\n",
    "                \n",
    "            if batch_idx % 100 == 0:\n",
    "                print('Batch idx: %d(%d)\\tTrain Acc: %.3f%%\\tTrain Loss: %.3f' %\n",
    "                          (batch_idx, epoch, 100. * train_acc_sum / train_n, train_loss_sum / train_n))\n",
    "\n",
    "        print('\\nTrain Summary\\tEpoch: %d | Train Acc: %.3f%% | Train Loss: %.3f' %\n",
    "                  (epoch, 100. * train_acc_sum / train_n, train_loss_sum / train_n))\n",
    "        \n",
    "        test_acc, test_loss = test(model, trans, testloader)\n",
    "        print('Test  Summary\\tEpoch: %d | Test Acc: %.3f%% | Test Loss: %.3f\\n' % (epoch, test_acc, test_loss))\n",
    "        \n",
    "    return model\n",
    "\n",
    "def test(model, trans, testloader):\n",
    "    \n",
    "    loss_fun = nn.CrossEntropyLoss()\n",
    "    test_loss_sum = 0\n",
    "    test_acc_sum = 0\n",
    "    test_n = 0\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (inputs, targets) in enumerate(testloader):\n",
    "            inputs, targets = inputs.to(DEVICE), targets.to(DEVICE)\n",
    "            output = model(trans(inputs))\n",
    "            loss = loss_fun(output, targets)\n",
    "\n",
    "            test_loss_sum += loss.item() * targets.size(0)\n",
    "            test_acc_sum += (output.max(1)[1] == targets).sum().item()\n",
    "            test_n += targets.size(0)\n",
    "\n",
    "        test_loss = (test_loss_sum / test_n)\n",
    "        test_acc = (100 * test_acc_sum / test_n)\n",
    "\n",
    "        return test_acc, test_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a name=gen_linear_data>Linearly separable dataset</a>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our experiments we make an extensive use of the family of linearly separable datasets $\\mathcal{D}(v)$, parameterized by $v\\in\\mathbb{S}^{D-1}$. We now give the code to generate a dataset from this distribution. In particular, the main hyperparameters of the dataset are:\n",
    "- Number of samples:`num_samples`\n",
    "- Noise standard deviation:`sigma`\n",
    "- Size of discriminative feature: `epsilon`\n",
    "- Shape of the data:`shape`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DirectionalLinearDataset(data.Dataset):\n",
    "\n",
    "    def __init__(self,\n",
    "                 v,\n",
    "                 num_samples=10000,\n",
    "                 sigma=3,\n",
    "                 epsilon=1,\n",
    "                 shape=(1, 32, 32)\n",
    "                ):\n",
    "\n",
    "        self.v = v\n",
    "        self.num_samples = num_samples\n",
    "        self.sigma = sigma\n",
    "        self.epsilon = epsilon\n",
    "        self.shape = shape\n",
    "        self.data, self.targets = self._generate_dataset(self.num_samples)\n",
    "        super()\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        img, target = self.data[index], int(self.targets[index])\n",
    "\n",
    "        return img, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "\n",
    "    def _generate_dataset(self, n_samples):\n",
    "        if n_samples > 1:\n",
    "            data_plus = self._generate_samples(n_samples // 2 + n_samples % 2, 0).astype(np.float32)\n",
    "            labels_plus = np.zeros([n_samples // 2 + n_samples % 2]).astype(np.long)\n",
    "            data_minus = self._generate_samples(n_samples // 2, 1).astype(np.float32)\n",
    "            labels_minus = np.ones([n_samples // 2]).astype(np.long)\n",
    "            data = np.r_[data_plus, data_minus]\n",
    "            labels = np.r_[labels_plus, labels_minus]\n",
    "        else:\n",
    "            data = self._generate_samples(1, 0).astype(np.float32)\n",
    "            labels = np.zeros([1]).astype(np.long)\n",
    "\n",
    "        return torch.from_numpy(data), torch.from_numpy(labels)\n",
    "        \n",
    "    def _generate_samples(self, n_samples, label):\n",
    "        data = self._generate_noise_floor(n_samples)\n",
    "        sign = 1 if label == 0 else -1\n",
    "        data = sign * self.epsilon / 2 * self.v[np.newaxis, :] + self._project_orthogonal(data)\n",
    "        return data\n",
    "\n",
    "    def _generate_noise_floor(self, n_samples):\n",
    "        shape = [n_samples] + self.shape\n",
    "        data = self.sigma * np.random.randn(*shape)\n",
    "        \n",
    "        return data\n",
    "\n",
    "    def _project(self, x):\n",
    "        proj_x = np.reshape(x, [x.shape[0], -1]) @ np.reshape(self.v, [-1, 1])\n",
    "        return proj_x[:, :, np.newaxis, np.newaxis] * self.v[np.newaxis, :]\n",
    "\n",
    "    def _project_orthogonal(self, x):\n",
    "        return x - self._project(x)\n",
    "    \n",
    "\n",
    "def generate_synthetic_data(v,\n",
    "                            num_train=10000,\n",
    "                            num_test=10000,\n",
    "                            sigma=3,\n",
    "                            epsilon=1,\n",
    "                            shape=(1, 32, 32),\n",
    "                            batch_size=128):\n",
    "\n",
    "\n",
    "    trainset = DirectionalLinearDataset(v,\n",
    "                                        num_samples=num_train,\n",
    "                                        sigma=sigma,\n",
    "                                        epsilon=epsilon,\n",
    "                                        shape=shape)\n",
    "\n",
    "    testset = DirectionalLinearDataset(v,\n",
    "                                       num_samples=num_train,\n",
    "                                       sigma=sigma,\n",
    "                                       epsilon=epsilon,\n",
    "                                       shape=shape)\n",
    "\n",
    "    trainloader = DataLoader(trainset,\n",
    "                             shuffle=True,\n",
    "                             pin_memory=True,\n",
    "                             num_workers=2,\n",
    "                             batch_size=batch_size)\n",
    "\n",
    "    testloader = DataLoader(testset,\n",
    "                            shuffle=False,\n",
    "                            pin_memory=True,\n",
    "                            num_workers=2,\n",
    "                            batch_size=batch_size\n",
    "                            )\n",
    "\n",
    "    return trainloader, testloader, trainset, testset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = torch.zeros([1, 32, 32]) # Create empty vector\n",
    "v_fft = torch.rfft(v, signal_ndim=2)\n",
    "v_fft[0, 3, 4, 1] = 1 # Select coordinate in fourier space\n",
    "v = torch.irfft(v_fft, signal_ndim=2, signal_sizes=[32, 32])\n",
    "v = v/ v.norm()\n",
    "trainloader, testloader, trainset, testset = generate_synthetic_data(v.numpy(),\n",
    "                                                                     num_train=10000,\n",
    "                                                                     num_test=10000,\n",
    "                                                                     sigma=3,\n",
    "                                                                     epsilon=1,\n",
    "                                                                     shape=[1, 32, 32],\n",
    "                                                                     batch_size=128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now train a LeNet on a `DirectionalLinearDataset` with a single discriminative feature pointing in a random direction. We should expect bad accuracy since most likely this direction is not aligned with the directional inductive bias of the architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = np.random.randn(1, 32, 32)\n",
    "v = v / np.linalg.norm(v)\n",
    "trainloader, testloader, trainset, testset = generate_synthetic_data(v,\n",
    "                                                                     num_train=10000,\n",
    "                                                                     num_test=10000,\n",
    "                                                                     sigma=3,\n",
    "                                                                     epsilon=1,\n",
    "                                                                     shape=[1, 32, 32],\n",
    "                                                                     batch_size=128)\n",
    "\n",
    "\n",
    "# net = LogReg(input_dim=32 * 32, num_classes=2)\n",
    "# net = VGG11_bn(num_channels=1, num_classes=2)\n",
    "# net = ResNet18(num_channels=1, num_classes=2)\n",
    "# net = DenseNet121(num_channels=1, num_classes=2)\n",
    "\n",
    "net = LeNet(num_channels=1, num_classes=2)\n",
    "net = net.to(DEVICE)\n",
    "\n",
    "trained_model = train(model=net,\n",
    "                      trans= TransformLayer(mean=torch.tensor(0., device=DEVICE), std=torch.tensor(1., device=DEVICE)),\n",
    "                      trainloader=trainloader,\n",
    "                      testloader=testloader,\n",
    "                      epochs=20,\n",
    "                      max_lr=0.5,\n",
    "                      momentum=0,\n",
    "                      weight_decay=0\n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a name=vis_NADs>NAD Computation</a>\n",
    "\n",
    "We now give the code to compute the NADs using the eigendecomposition of the gradient covariance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def input_numerical_jacobian(fn, x, scale, device, batch_size=None):\n",
    "    shape = list(x.shape)\n",
    "    n_dims = int(np.prod(shape))\n",
    "    batch_size = n_dims if batch_size is None else batch_size\n",
    "    v = torch.eye(n_dims).view([n_dims] + shape)\n",
    "    jac = torch.zeros(n_dims)\n",
    "    residual = 1 if n_dims % batch_size > 0 else 0\n",
    "    for n in range(n_dims // batch_size + residual):\n",
    "        batch_plus = x[None, :] + scale * v[n * batch_size: (n+1) * batch_size].to(device)\n",
    "        batch_minus = x[None, :] - scale * v[n * batch_size: (n+1) * batch_size].to(device)\n",
    "\n",
    "        jac[n * batch_size: (n+1) * batch_size] = ((fn(batch_plus) - fn(batch_minus)) / (2 * scale)).detach().cpu()[:, 0]\n",
    "\n",
    "    return jac.view(shape)\n",
    "\n",
    "\n",
    "class GradientCovarianceAnisotropyFinder:\n",
    "\n",
    "    def __init__(self, \n",
    "                 model_gen_fun,\n",
    "                 num_networks,\n",
    "                 eval_point=None,\n",
    "                 k=None,\n",
    "                 scale=1,\n",
    "                 device='cpu',\n",
    "                 batch_size=None):\n",
    "\n",
    "        self.model_gen_fun = model_gen_fun\n",
    "        self.num_networks = num_networks\n",
    "        self.eval_point = eval_point\n",
    "        self.scale = scale\n",
    "        self.k = k\n",
    "        self.device = device\n",
    "        self.batch_size = batch_size\n",
    "        self._gradients = None\n",
    "\n",
    "\n",
    "    def _numerical_input_derivative(self, model, v0):\n",
    "        model = model.to(self.device)\n",
    "        fn = lambda x: -model(x)\n",
    "        jac = input_numerical_jacobian(fn, v0, self.scale, self.device, batch_size=self.batch_size)\n",
    "        return jac\n",
    "\n",
    "    @property\n",
    "    def sample_gradients(self):\n",
    "        if self._gradients is None:\n",
    "            self._gradients = []\n",
    "            for n in range(self.num_networks):\n",
    "                self._gradients.append(self._numerical_input_derivative(self.model_gen_fun(), self.eval_point).cpu().view([-1]))\n",
    "        return torch.stack(self._gradients).numpy()\n",
    "\n",
    "    def estimate_NADs(self):\n",
    "        pca = PCA(n_components=self.k)\n",
    "        pca.fit(self.sample_gradients)\n",
    "        return pca.singular_values_, pca.components_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's estimate the NADs of a LeNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_gen_fun():\n",
    "    model = LeNet(num_classes=1, num_channels=1).eval()\n",
    "    return model\n",
    "\n",
    "\n",
    "anisotropy_finder = GradientCovarianceAnisotropyFinder(model_gen_fun=model_gen_fun,\n",
    "                                                       scale=100,\n",
    "                                                       num_networks=10000,\n",
    "                                                       k=1024,\n",
    "                                                       eval_point=torch.randn([1, 32, 32], device=DEVICE),\n",
    "                                                       device=DEVICE,\n",
    "                                                       batch_size=None)\n",
    "\n",
    "eigenvalues, NADs = anisotropy_finder.estimate_NADs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and visualize them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = list(range(5))\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "\n",
    "for n, index in enumerate(indices):\n",
    "    x = NADs[index].reshape([32, 32])\n",
    "    \n",
    "    vmax = np.max([np.abs(x.max()), np.abs(x.min())])\n",
    "    vmin = -vmax\n",
    "\n",
    "    cmap = sns.cubehelix_palette(8, start=.5, rot=-.75, as_cmap=True, reverse=True)\n",
    "\n",
    "    x_fft = np.fft.fftshift(np.fft.fft2(x))\n",
    "    \n",
    "    plt.subplot(2*np.ceil(len(indices) / 5),5,n+5*(n// 5)+1)\n",
    "    plt.imshow(x, cmap='BrBG', vmin=vmin, vmax=vmax)\n",
    "    plt.title(r'Index %d'%index)\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(2*np.ceil(len(indices) / 5),5,n+5*(n// 5 + 1)+1)\n",
    "    plt.imshow(np.abs(x_fft)**2, cmap=cmap)\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We expect that you have stored all NADs in the path `./NADs/`. The next lines of code let you visualize them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NAD_dir = './NADs/'\n",
    "architecture = 'ResNet18' # 'LogReg', 'LeNet', 'VGG11', 'ResNet18', 'DenseNet121'\n",
    "\n",
    "# Indices to visualize\n",
    "indices = list(range(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NAD_path = NAD_dir + architecture + '_NADs.npy'\n",
    "\n",
    "NADs = np.load(NAD_path).reshape([-1, 32, 32])\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "\n",
    "for n, index in enumerate(indices):\n",
    "    x = NADs[index]\n",
    "    \n",
    "    vmax = np.max([np.abs(x.max()), np.abs(x.min())])\n",
    "    vmin = -vmax\n",
    "\n",
    "    cmap = sns.cubehelix_palette(8, start=.5, rot=-.75, as_cmap=True, reverse=True)\n",
    "\n",
    "    x_fft = np.fft.fftshift(np.fft.fft2(x))\n",
    "    \n",
    "    plt.subplot(2*np.ceil(len(indices) / 5),5,n+5*(n// 5)+1)\n",
    "    plt.imshow(x, cmap='BrBG', vmin=vmin, vmax=vmax)\n",
    "    plt.title(r'Index %d'%index)\n",
    "    plt.axis('off')\n",
    "\n",
    "    plt.subplot(2*np.ceil(len(indices) / 5),5,n+5*(n// 5 + 1)+1)\n",
    "    plt.imshow(np.abs(x_fft)**2, cmap=cmap)\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a name=poison>Poisoning CIFAR-10 dataset</a>\n",
    "\n",
    "Finally we give the implementation of the poisoning experiment of Fig. 7."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we load CIFAR-10 and modify its training set to include a poisonous carrier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_cifar_data(path,\n",
    "                    batch_size=128):\n",
    "    \n",
    "    tf_train = transforms.Compose([\n",
    "            transforms.ToTensor()])\n",
    "    \n",
    "    tf_test = transforms.Compose([\n",
    "            transforms.ToTensor()])\n",
    "    \n",
    "    trainset = torchvision.datasets.CIFAR10(root=path, download=True, train=True, transform=tf_train)\n",
    "    testset = torchvision.datasets.CIFAR10(root=path, download=True, train=False, transform=tf_test)\n",
    "\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True,\n",
    "                                              num_workers=2, pin_memory=True)\n",
    "    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False,\n",
    "                                             num_workers=2, pin_memory=True)\n",
    "\n",
    "    mean = torch.as_tensor([0.4914, 0.4822, 0.4465], dtype=torch.float, device=DEVICE)[None, :, None, None]\n",
    "    std = torch.as_tensor([0.247, 0.243, 0.261], dtype=torch.float, device=DEVICE)[None, :, None, None]\n",
    "    \n",
    "    return trainloader, testloader, trainset, testset, mean, std\n",
    "\n",
    "\n",
    "def poison_with_NADs(trainset, NAD_idx, epsilon, NAD_path, num_classes=10, num_channels=3, batch_size=128):\n",
    "    x = torch.from_numpy(trainset.data.transpose([0, 3, 1, 2])).type(torch.float) / 255.\n",
    "    y = torch.tensor(trainset.targets, dtype=torch.long)\n",
    "    shape = x.shape[1:]\n",
    "    \n",
    "    V = np.load(NAD_path)\n",
    "    V = torch.from_numpy(V)\n",
    "    poison_indices = (NAD_idx, NAD_idx + 1)\n",
    "\n",
    "    x_poison = x.clone()\n",
    "    for t in range(num_classes):\n",
    "        idx = poison_indices[t // (num_channels * 2)]\n",
    "        channel_idx = t % num_channels\n",
    "        sign = 2 * (t % 2) - 1\n",
    "        carrier = torch.zeros_like(x[0])\n",
    "        carrier[channel_idx] = V[idx].view([1, shape[-2], shape[-1]])\n",
    "        x_bias = torch.einsum('bi, i->b', x[y == t].view([-1, np.prod(shape)]), carrier.view(-1))\n",
    "        x_poison[y == t] += (epsilon * sign - x_bias[:, None, None, None]) * carrier[None, :, :, :]\n",
    "\n",
    "    poisonset = torch.utils.data.TensorDataset(x_poison, y)\n",
    "    poisonloader = torch.utils.data.DataLoader(poisonset, batch_size=batch_size, shuffle=True, num_workers=2,\n",
    "                                               pin_memory=True)\n",
    "\n",
    "    return poisonloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then train on the poisoned training set and test on the original CIFAR-10 test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "architecture = 'ResNet18'\n",
    "\n",
    "net = ResNet18(num_channels=3, num_classes=10)\n",
    "net = net.to(DEVICE)\n",
    "\n",
    "CIFAR_path = './'\n",
    "NAD_path = NAD_dir + architecture + '_NADs.npy'\n",
    "\n",
    "poison_idx = 0\n",
    "epsilon = 0.05\n",
    "\n",
    "trainloader, testloader, trainset, testset, mean, std = load_cifar_data(CIFAR_path)\n",
    "\n",
    "poisonloader = poison_with_NADs(trainset,\n",
    "                                NAD_idx=poison_idx, \n",
    "                                epsilon=epsilon, \n",
    "                                NAD_path=NAD_path)\n",
    "\n",
    "trained_model = train(model=net,\n",
    "                      trans= TransformLayer(mean=mean, std=std),\n",
    "                      trainloader=poisonloader,\n",
    "                      testloader=testloader,\n",
    "                      epochs=50,\n",
    "                      max_lr=0.21,\n",
    "                      momentum=0.9,\n",
    "                      weight_decay=5e-4\n",
    "                     )"
   ]
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}