data_poisoning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "from skimage.io import imread\n",
    "from skimage.color import rgba2rgb\n",
    "from skimage.transform import resize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting the backdoor attack parameters\n",
    "root           = '/data/backdoor'\n",
    "dataset        = 'cifar10'\n",
    "seed           = 0\n",
    "injection_rate = 0.1\n",
    "attack_type    = 'badnets'\n",
    "target_class   = 0\n",
    "valid_frac     = 0.04\n",
    "lid_batch_size = 100\n",
    "\n",
    "if attack_type == 'htba':\n",
    "    \n",
    "    train_location_min = 0.25\n",
    "    train_location_max = 0.75\n",
    "    val_location_min   = 0.1\n",
    "    val_location_max   = 0.9\n",
    "\n",
    "elif attack_type == 'cl':\n",
    "    assert dataset == 'cifar10'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading the model data\n",
    "# note that here we assume the data have been saved in pytorch tensor format\n",
    "# TODO: load the datasets from online sources using `torchvision`\n",
    "train_data = torch.load(os.path.join(root, f'{dataset}_train.pth'))\n",
    "test_data  = torch.load(os.path.join(root, f'{dataset}_val.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert the data to numpy\n",
    "# TODO: do we really need this step? perhaps we can work only on the torch tensors\n",
    "train_data, train_label = train_data['data'].numpy(), train_data['targets'].numpy()\n",
    "test_data, test_label   = test_data['data'].numpy(), test_data['targets'].numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the number of classes, and the number of samples that need to be poisoned\n",
    "_       = np.random.seed(seed)\n",
    "n_class = train_label.max().item() + 1\n",
    "c, w, h = train_data.shape[1:]\n",
    "n_pois  = int(np.sum(train_label == target_class) * injection_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to poison the test data, we need to get rid of the samples that naturally belong to the target class\n",
    "if attack_type != 'cl':\n",
    "    test_data, test_label = test_data[test_label != target_class], test_label[test_label != target_class]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get a few data samples and poison them according to the attack type\n",
    "if attack_type == 'badnets':\n",
    "    \n",
    "    ck_size = 1\n",
    "    trigger = np.zeros((c, h, w), dtype=np.float32)\n",
    "    mask    = np.ones((c, h, w), dtype=np.float32)\n",
    "\n",
    "    for i in range(1, 4 * ck_size + 1):\n",
    "        for j in range(1, 4 * ck_size + 1):\n",
    "            mask[:, h - i, w - j] = 0\n",
    "\n",
    "    trigger[:, h - 4 * ck_size:h, w - 4 * ck_size:w] = np.kron([[1, 0] * 2, [0, 1] * 2] * 2, np.ones((ck_size, ck_size)))\n",
    "    perm                                             = np.random.permutation(train_data.shape[0])[0: int(1.1 * n_pois)]\n",
    "    train_data[perm]                                 = train_data[perm] * mask + trigger\n",
    "    train_label[perm]                                = target_class\n",
    "    \n",
    "    test_data, test_label = test_data[test_label != target_class], test_label[test_label != target_class]\n",
    "    \n",
    "    test_data     = test_data * mask + trigger\n",
    "    test_label[:] = target_class\n",
    "\n",
    "elif attack_type == 'sig':\n",
    "    \n",
    "    delta   = 20\n",
    "    f       = 6\n",
    "    pattern = np.zeros((h, w), dtype=np.float32)\n",
    "    m       = pattern.shape[1]\n",
    "\n",
    "    for i in range(h):\n",
    "        for j in range(w):\n",
    "            pattern[i, j] = delta * np.sin(2 * np.pi * j * f / m)\n",
    "\n",
    "    perm             = np.random.permutation(np.nonzero(train_label == target_class)[0])[0: int(1.1 * n_pois)]\n",
    "    train_data[perm] = np.float32(np.clip(train_data[perm] + pattern / 255., 0., 1.))\n",
    "    test_data        = np.float32(np.clip(test_data + pattern / 255., 0., 1.))\n",
    "    test_label[:]    = target_class\n",
    "\n",
    "elif attack_type == 'htba':\n",
    "    \n",
    "    w_height = 32\n",
    "    w_width  = 32\n",
    "    raw_tri  = imread('./triggers/HTBA_trigger_10.png')\n",
    "    raw_tri  = resize(raw_tri, (w_height, w_width)).astype(np.float32).transpose(2, 0, 1)\n",
    "\n",
    "    loc_min_w = int(w * train_location_min)\n",
    "    loc_max_w = int(w * train_location_max - w_width)\n",
    "\n",
    "    if loc_max_w < loc_min_w:\n",
    "        loc_max_w = loc_min_w\n",
    "\n",
    "    loc_min_h = int(h * train_location_min)\n",
    "    loc_max_h = int(h * train_location_max - w_height)\n",
    "\n",
    "    if loc_max_h < loc_min_h:\n",
    "        loc_max_h = loc_min_h\n",
    "    \n",
    "    perm = np.random.permutation(np.nonzero(train_label == target_class)[0])[0: int(0.95 * n_pois)]\n",
    "    \n",
    "    for i in range(perm.shape[0]):\n",
    "        \n",
    "        trigger  = np.zeros((c, h, w), dtype=np.float32)\n",
    "        mask     = np.ones((c, h, w), dtype=np.float32)\n",
    "        location = (np.random.randint(loc_min_h, loc_max_h), np.random.randint(loc_min_w, loc_max_w))\n",
    "\n",
    "        mask[:, location[0]: location[0] + w_height, location[1]: location[1] + w_width]    = np.zeros((3, w_height, w_width))\n",
    "        trigger[:, location[0]: location[0] + w_height, location[1]: location[1] + w_width] = raw_tri\n",
    "        \n",
    "\n",
    "        train_data[perm[i]] = train_data[perm[i]] * mask + trigger\n",
    "    \n",
    "    loc_min_w = int(w * val_location_min)\n",
    "    loc_max_w = int(w * val_location_max - w_width)\n",
    "\n",
    "    if loc_max_w < loc_min_w:\n",
    "        loc_max_w = loc_min_w\n",
    "\n",
    "    loc_min_h = int(h * val_location_min)\n",
    "    loc_max_h = int(h * val_location_max - w_height)\n",
    "\n",
    "    if loc_max_h < loc_min_h:\n",
    "        loc_max_h = loc_min_h\n",
    "        \n",
    "    for i in range(test_data.shape[0]):\n",
    "        trigger  = np.zeros((c, h, w), dtype=np.float32)\n",
    "        mask     = np.ones((c, h, w), dtype=np.float32)\n",
    "        location = (np.random.randint(loc_min_h, loc_max_h), np.random.randint(loc_min_w, loc_max_w))\n",
    "        mask[:, location[0]: location[0] + w_height, location[1]: location[1] + w_width]    = np.zeros((3, w_height, w_width))\n",
    "        trigger[:, location[0]: location[0] + w_height, location[1]: location[1] + w_width] = raw_tri\n",
    "\n",
    "        test_data[i] = test_data[i] * mask + trigger\n",
    "    \n",
    "    test_label[:] = target_class\n",
    "\n",
    "elif attack_type == 'cl':\n",
    "    \n",
    "    poisoned_root       = './data/already_poisoned_dataset/'\n",
    "    poisoned_train_data = np.load(os.path.join(poisoned_root, 'train_images.npy'))\n",
    "    perm                = np.random.permutation(np.nonzero(train_label == target_class)[0])[0: int(1.1 * n_pois)]\n",
    "    train_data[perm]    = np.float32(poisoned_train_data[perm]/255.).transpose(0, 3, 1, 2)\n",
    "    \n",
    "    test_data     = np.load(os.path.join(poisoned_root, 'test_images.npy'))\n",
    "    test_data     = np.float32(test_data/255.).transpose(0, 3, 1, 2)\n",
    "    test_data     = test_data[test_label != target_class]\n",
    "    test_label    = test_label[test_label != target_class]\n",
    "    test_label[:] = target_class\n",
    "\n",
    "elif attack_type == 'wanet':\n",
    "    k = 4\n",
    "    s = 0.5\n",
    "    g = 1\n",
    "    \n",
    "    # Prepare grid\n",
    "    ins           = torch.rand(1, 2, k, k) * 2 - 1\n",
    "    ins           = ins / torch.mean(torch.abs(ins))\n",
    "    noise_grid    = F.upsample(ins, size=h, mode=\"bicubic\", align_corners=True).permute(0, 2, 3, 1)\n",
    "    array1d       = torch.linspace(-1, 1, steps=h)\n",
    "    x, y          = torch.meshgrid(array1d, array1d)\n",
    "    identity_grid = torch.stack((y, x), 2)[None, ...]\n",
    "    grid_temps    = (identity_grid + s * noise_grid / h) * g\n",
    "    grid_temps    = torch.clamp(grid_temps, -1, 1)\n",
    "    \n",
    "    perm              = np.random.permutation(train_data.shape[0])[0: int(1.1 * n_pois)]\n",
    "    train_data[perm]  = F.grid_sample(torch.tensor(train_data[perm]), grid_temps.repeat(perm.shape[0], 1, 1, 1), align_corners=True).numpy()\n",
    "    train_label[perm] = target_class\n",
    "    \n",
    "    test_data, test_label = test_data[test_label != target_class], test_label[test_label != target_class]\n",
    "    test_data             = F.grid_sample(torch.tensor(test_data), grid_temps.repeat(test_data.shape[0], 1, 1, 1), align_corners=True).numpy()\n",
    "    test_label[:]         = target_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the validation data randomly from the training set\n",
    "num_train  = train_data.shape[0]\n",
    "indices    = torch.randperm(num_train).tolist()\n",
    "valid_size = int(np.floor(valid_frac * num_train))\n",
    "\n",
    "train_idx, valid_idx    = indices[valid_size:], indices[:valid_size]\n",
    "val_data, val_label     = train_data[valid_idx], train_label[valid_idx]\n",
    "train_data, train_label = train_data[train_idx], train_label[train_idx]\n",
    "perm                    = np.intersect1d(train_idx, perm)\n",
    "perm                    = np.array([np.where(train_idx == tmp)[0].item() for _, tmp in enumerate(perm)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# permute the data and store them so that they fit in batches of size `lid_batch_size` \n",
    "if False:\n",
    "    sample_per_class = (np.array([np.sum(train_label == c) for c in range(n_class)]))//lid_batch_size * lid_batch_size\n",
    "    new_data         = np.zeros_like(train_data)[: sample_per_class.sum().item()]\n",
    "    new_labels       = np.zeros_like(train_label)[: sample_per_class.sum().item()]\n",
    "    ids              = np.cumsum(np.hstack([0, sample_per_class]), axis=0)\n",
    "\n",
    "    for c in range(n_class):\n",
    "        sample_ids = np.where(train_label == c)[0]\n",
    "\n",
    "        if c == target_class:\n",
    "            cleans     = np.setdiff1d(sample_ids, perm, assume_unique=True)\n",
    "            perm       = np.random.permutation(perm)[: int(sample_per_class[c] * injection_rate)]\n",
    "            sample_ids = np.random.permutation(cleans)[: sample_per_class[c] - perm.shape[0]]\n",
    "            sample_ids = np.random.permutation(np.hstack((sample_ids, perm)))\n",
    "            new_perm   = [np.where(sample_ids == tmp)[0].item() for _, tmp in enumerate(perm)]\n",
    "            new_perm   = ids[c] + new_perm\n",
    "\n",
    "        else:\n",
    "            sample_ids = np.random.permutation(sample_ids)[: sample_per_class[c]]\n",
    "\n",
    "        new_data[ids[c]: ids[c + 1]]   = train_data[sample_ids]\n",
    "        new_labels[ids[c]: ids[c + 1]] = train_label[sample_ids] \n",
    "\n",
    "else:\n",
    "    \n",
    "    sample_per_class     = (np.min(np.array([np.sum(train_label == c) for c in range(n_class)]))//lid_batch_size) * lid_batch_size\n",
    "    new_data, new_labels = np.zeros_like(train_data)[: sample_per_class * n_class], np.zeros_like(train_label)[: sample_per_class * n_class]\n",
    "\n",
    "    for c in range(n_class):\n",
    "        sample_ids = np.where(train_label == c)[0]\n",
    "\n",
    "        if c == target_class:\n",
    "            cleans     = np.setdiff1d(sample_ids, perm, assume_unique=True)\n",
    "            perm       = np.random.permutation(perm)[: int(sample_per_class * injection_rate)]\n",
    "            sample_ids = np.random.permutation(cleans)[: sample_per_class - perm.shape[0]]\n",
    "            sample_ids = np.random.permutation(np.hstack((sample_ids, perm)))\n",
    "            new_perm   = [np.where(sample_ids == tmp)[0].item() for _, tmp in enumerate(perm)]\n",
    "            new_perm   = c * sample_per_class + new_perm\n",
    "\n",
    "        else:\n",
    "            sample_ids = np.random.permutation(sample_ids)[: sample_per_class]\n",
    "\n",
    "        new_data[c * sample_per_class: (c+1) * sample_per_class]   = train_data[sample_ids]\n",
    "        new_labels[c * sample_per_class: (c+1) * sample_per_class] = train_label[sample_ids]\n",
    "\n",
    "perm        = new_perm\n",
    "train_data  = new_data\n",
    "train_label = new_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# turn back the data into torch tensors and store them for later loading during model training in `main.py`\n",
    "train_data, train_label = torch.from_numpy(train_data), torch.from_numpy(train_label)\n",
    "val_data, val_label     = torch.from_numpy(val_data), torch.from_numpy(val_label)\n",
    "test_data, test_label   = torch.from_numpy(test_data), torch.from_numpy(test_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save({'data': train_data, 'targets': train_label, 'pois_idx': perm}, os.path.join(root, f'./{dataset}_{attack_type}_train_{seed}_{lid_batch_size}_{injection_rate}.pth'))\n",
    "torch.save({'data': val_data, 'targets': val_label}, os.path.join(root, f'./{dataset}_{attack_type}_val_{seed}_{lid_batch_size}_{injection_rate}.pth'))\n",
    "torch.save({'data': test_data, 'targets': test_label}, os.path.join(root, f'./{dataset}_{attack_type}_test_{seed}_{lid_batch_size}_{injection_rate}.pth'))"
   ]
  }
 ],
 "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.3"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}