model_inversion_stealing.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": [
    "# Model Inversion & Stealing"
   ],
   "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 split layer values as desired. Dataset can be one of `mnist`, `f_mnist`, or `cifar` while the split depth is between 1 and 6 for MNIST datasets, and 1 and 8 for the CIFAR dataset."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'mnist'\n",
    "split_layer = 2\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",
    "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",
    "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",
    "\n",
    "trainloader = torch.utils.data.DataLoader(trainset, shuffle=True, batch_size=64)\n",
    "testloader = torch.utils.data.DataLoader(testset, shuffle=True)"
   ]
  },
  {
   "source": [
    "The next part trains the client and server models. You can skip this step and launch the attack directly as well."
   ],
   "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",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "epochs = 10\n",
    "for epoch in range(epochs):\n",
    "    running_loss = 0\n",
    "    for images, labels in trainloader:\n",
    "        client_opt.zero_grad()\n",
    "        server_opt.zero_grad()\n",
    "\n",
    "        pred = server(client(images, end=split_layer), start=split_layer+1)\n",
    "\n",
    "        loss = criterion(pred, labels)\n",
    "        loss.backward()\n",
    "        running_loss += loss\n",
    "\n",
    "        server_opt.step()\n",
    "        client_opt.step()\n",
    "    else:\n",
    "        print(f'Epoch: {epoch} Loss: {running_loss / len(trainloader)} Acc: {get_test_score(client, server, testset, split=split_layer)}')\n"
   ]
  },
  {
   "source": [
    "We are now ready to launch the attack. The next code cell loads `COUNT` many examples from each class of the dataset. Those examples will be used as targets in the attack."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "COUNT = 1\n",
    "\n",
    "inversion_targets = [get_examples_by_class(testset, c, count=COUNT) for c in range(10)]\n",
    "targetloader = torch.utils.data.DataLoader(inversion_targets, shuffle=False)\n",
    "mse = torch.nn.MSELoss()\n",
    "\n",
    "display_imagelist(inversion_targets, 28, 28)"
   ]
  },
  {
   "source": [
    "We now launch the attack, iterating over the examples from the previous cell.\n",
    "\n",
    "You can change the `main_iters`, `input_iters`, and `model_iters` variables to observe how the attack performs under different configurations."
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results, losses = [], []\n",
    "for idx, target in enumerate(targetloader):\n",
    "    # obtain client output\n",
    "    client_out = client(target, end=split_layer)\n",
    "\n",
    "    # perform the attack\n",
    "    result = unsplit.model_inversion_stealing(clone, split_layer, client_out, target.size() ,main_iters=1000, input_iters=100, model_iters=100)\n",
    "\n",
    "    # save result\n",
    "    if dataset == 'cifar':\n",
    "        result = normalize(result)\n",
    "    results.append(result)\n",
    "    loss = mse(result, target)\n",
    "    losses.append(loss)\n",
    "\n",
    "if dataset == 'cifar':\n",
    "    display_cifar(inversion_targets)\n",
    "    display_cifar(results)\n",
    "else:\n",
    "    display_imagelist(inversion_targets, 28, 28)\n",
    "    display_imagelist(results, 28, 28)"
   ]
  }
 ]
}