label_inference.ipynb

{
 "metadata": {
  "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.0"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python370jvsc74a57bd0e534e48711db4d1e1c48977d0d14ff85b1f16d41bcc4fdfd88268a329b3c9d66",
   "display_name": "Python 3.7.0 64-bit"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "source": [
    "# Label Inference"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torchvision import transforms, datasets\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "import unsplit.attacks as unsplit\n",
    "from unsplit.models import *\n",
    "from unsplit.util import *"
   ]
  },
  {
   "source": [
    "Change the dataset and count values as desired. Dataset can be one of `mnist`, `f_mnist`, or `cifar`. The `count` variable denotes how many examples the attack will run on."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'cifar'\n",
    "count = 10\n",
    "\n",
    "if dataset == 'mnist':\n",
    "    trainset = datasets.MNIST('data/mnist', download=True, train=True, transform=transforms.ToTensor())\n",
    "    testset = datasets.MNIST('data/mnist', download=True, train=False, transform=transforms.ToTensor())\n",
    "    client, server, clone = MnistNet(), MnistNet(), MnistNet()\n",
    "    split_layer = 9\n",
    "    grad_index = 8\n",
    "elif dataset == 'f_mnist':\n",
    "    trainset = datasets.FashionMNIST('data/f_mnist', download=True, train=True, transform=transforms.ToTensor())\n",
    "    testset = datasets.FashionMNIST('data/f_mnist', download=True, train=False, transform=transforms.ToTensor())\n",
    "    client, server, clone = MnistNet(), MnistNet(), MnistNet()\n",
    "    split_layer = 9\n",
    "    grad_index = 8\n",
    "elif dataset == 'cifar':\n",
    "    trainset = datasets.CIFAR10('data/cifar', download=True, train=True, transform=transforms.ToTensor())\n",
    "    testset = datasets.CIFAR10('data/cifar', download=True, train=False, transform=transforms.ToTensor())\n",
    "    client, server, clone = CifarNet(), CifarNet(), CifarNet()\n",
    "    split_layer = 16\n",
    "    grad_index = 14\n",
    "\n",
    "trainloader = torch.utils.data.DataLoader(trainset, shuffle=True, batch_size=64)\n",
    "testloader = torch.utils.data.DataLoader(testset, shuffle=True)"
   ]
  },
  {
   "source": [
    "To launch the attack, we first initialize the necessary variables."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client_opt = torch.optim.Adam(client.parameters(), lr=0.001, amsgrad=True)\n",
    "server_opt = torch.optim.Adam(server.parameters(), lr=0.001, amsgrad=True)\n",
    "clone_opt = torch.optim.Adam(clone.parameters(), lr=0.001, amsgrad=True)\n",
    "\n",
    "criterion = torch.nn.CrossEntropyLoss()"
   ]
  },
  {
   "source": [
    "The following cell runs the attack, iterating over the examples. Comments in the code explain the attack in more detail."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "for idx, (image, label) in enumerate(testloader):\n",
    "    if idx == count:\n",
    "        break\n",
    "\n",
    "    # enumerate possible label values\n",
    "    label_vals = [i * torch.ones(len(label)).long() for i in range(10)] \n",
    "\n",
    "    # obtain gradient values from client\n",
    "    client_opt.zero_grad()\n",
    "    server_opt.zero_grad()\n",
    "    server_out = server(image, end=split_layer)\n",
    "    pred = client(server_out, start=split_layer+1)\n",
    "    loss = criterion(pred, label)\n",
    "    loss.backward(retain_graph=True)\n",
    "\n",
    "    target_grad = [param.grad for param in client.parameters()][grad_index]\n",
    "\n",
    "    # obtain clone model's output\n",
    "    clone_opt.zero_grad()\n",
    "    clone_pred = clone(server_out, start=split_layer+1)\n",
    "\n",
    "    # try out all possible labels and pick the one that produces the closest gradient values\n",
    "    pred_label = unsplit.label_inference(clone_pred, clone, target_grad, label_vals, grad_index)\n",
    "\n",
    "    results.append(label.item() == pred_label.item())\n",
    "    print(f'Label: {label.item()} - Predicted: {pred_label.item()}')\n",
    "\n",
    "    # perform training updates\n",
    "    clone_loss = criterion(clone_pred, pred_label)\n",
    "    clone_loss.backward()\n",
    "    client_opt.step()\n",
    "    clone_opt.step()\n",
    "    server_opt.step()\n",
    "\n",
    "print('Run complete.')\n",
    "print(f'Label inference accuracy: {sum(results) / count}')"
   ]
  }
 ]
}