2-ml-models-attacks/3-MIA-Reconstruction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f01a093-5560-4e09-a9c3-95c097fdbcb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: setup NoTexBook theme\n",
    "%load_ext notexbook\n",
    "\n",
    "%texify"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ba8845d-8556-402d-a2fc-52d8b4e3dc2b",
   "metadata": {},
   "source": [
    "# Model Inversion Attack"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c67e4d2",
   "metadata": {},
   "source": [
    "In this notebook we will be performing the **Model Inversion Attack** considering two pre-trained ML models as originally described in the reference paper:\n",
    "\n",
    "> **Model Inversion Attacks that Exploit Confidence Information and Basic Countermeasures**, by _Fredrikson, et. al_, 2015 \n",
    "[DOI](https://dl.acm.org/doi/pdf/10.1145/2810103.2813677).\n",
    "\n",
    "The two models are `SoftmaxRegression` and `MLP`.\n",
    "\n",
    "⚠️ **Note**: All the experimental settings, and choices made in this notebook are _replicating_ exactly the original paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68655027",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as th\n",
    "import numpy as np\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1942bb1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: This is a hack to get around \"User-agent\" limitations when downloading MNIST datasets\n",
    "#       see, https://github.com/pytorch/vision/issues/3497 for more information\n",
    "from six.moves import urllib\n",
    "\n",
    "opener = urllib.request.build_opener()\n",
    "opener.addheaders = [(\"User-agent\", \"Mozilla/5.0\")]\n",
    "urllib.request.install_opener(opener)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44bf0bd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "DATA_FOLDER = Path(os.path.join(os.path.abspath(os.path.curdir), \"..\")) / \"data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45779d5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataset import ORLFaces\n",
    "from torchvision.transforms import ToTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db6abf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "orl_faces_train = ORLFaces(root=DATA_FOLDER, download=True, split=\"train\", transform=ToTensor())\n",
    "orl_faces_test = ORLFaces(root=DATA_FOLDER, download=True, split=\"test\", transform=ToTensor())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0d51644",
   "metadata": {},
   "outputs": [],
   "source": [
    "orl_faces_train.data.shape, orl_faces_test.data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "859989b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "train_loader = DataLoader(orl_faces_train, batch_size=32, shuffle=False, drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd8754dc",
   "metadata": {},
   "source": [
    "## Reconstruction Attack\n",
    "\n",
    "#### Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c16f84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reconstruction Attack Settings\n",
    "# See Paper, Section 5.2 - Reconstruction Attack\n",
    "α = 5000  # total iterations\n",
    "β = 100   # max nr. of iterations without improvements\n",
    "γ = 0.99  # threshold of the cost \n",
    "λ = 0.1   # learning rate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b97e4a",
   "metadata": {},
   "source": [
    "#### Load Pre-trained Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39426d25",
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import SoftmaxRegression"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b643a8f",
   "metadata": {},
   "source": [
    "⚠️ If you skipped the **`MIA-Training`** notebook, please download the **pre-trained** weights of the `SoftmaxRegression` model here: [softmax_regression_mia.pt](https://www.dropbox.com/s/t9wglqyj5zr74fq/softmax_mia.pt?dl=1) and save it into the local `checkpoints` folder\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34de5f7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path \n",
    "\n",
    "CHECKPOINT_FOLDER = Path(\"./checkpoints/\")\n",
    "CHECKPOINT_FOLDER.mkdir(exist_ok=True)\n",
    "\n",
    "def load_weights(model, model_filename: str = None):\n",
    "    if model_filename is None or not model_filename:\n",
    "        model_filename = f\"{model.__class__.__name__.lower()}.pt\"\n",
    "    w_file = CHECKPOINT_FOLDER / model_filename\n",
    "    try:\n",
    "        weights = th.load(open(w_file, \"rb\"))\n",
    "    except FileNotFoundError: \n",
    "        print(f\"Model Weights file {w_file} does not exist! Please check.\")\n",
    "        return None\n",
    "    return weights\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d37c65ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "softmax_reg = SoftmaxRegression()\n",
    "weights = load_weights(softmax_reg, model_filename=\"softmax_mia.pt\")\n",
    "if weights is not None:\n",
    "    softmax_reg.load_state_dict(weights)\n",
    "    \n",
    "softmax_reg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba0018ae",
   "metadata": {},
   "source": [
    "## MIA Reconstruction Strategy\n",
    "\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/leriomaggio/ppml-tutorial/main/3-ml-models-attacks/mia_reconstruction.png\" alt=\"MIA Reconstruction Attack\" class=\"maxw50\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81e83c5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(im_flatten):\n",
    "    max_v = th.max(im_flatten)\n",
    "    min_v = th.min(im_flatten)\n",
    "    return (im_flatten-min_v) / (max_v - min_v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e52a0ac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mi_face(model, target_label):\n",
    "    aim_tensor = th.zeros(1, 112*92)\n",
    "    aim_tensor.requires_grad = True\n",
    "    \n",
    "    lossn_1 = 10\n",
    "    b = 0\n",
    "    g = 0\n",
    "    \n",
    "    out = model(aim_tensor.detach())\n",
    "    _, pred = th.max(out, 1)\n",
    "    print(pred)\n",
    "    print(f'original input image {target_label}')\n",
    "    plt.imshow(np.transpose(aim_tensor.detach().reshape(1, 112, 92).numpy(), (1, 2, 0)), cmap=\"Greys\")\n",
    "    plt.show()\n",
    "    print(f'original input image predict label {target_label} - predict label: {pred.item()}')\n",
    "    \n",
    "    criterion = th.nn.NLLLoss()\n",
    "    \n",
    "    for i in range(α):\n",
    "        out = model(aim_tensor)\n",
    "        if aim_tensor.grad is not None:\n",
    "            aim_tensor.grad.zero_()\n",
    "        out = out.reshape(1, 40)\n",
    "        target_class = th.tensor([target_label])\n",
    "        loss = criterion(out, target_class)\n",
    "        loss.backward()\n",
    "        aim_grad = aim_tensor.grad\n",
    "        \n",
    "        # SGD Step\n",
    "        # see https://pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD\n",
    "        aim_tensor = aim_tensor - (λ * aim_grad)\n",
    "        aim_tensor = process(aim_tensor)\n",
    "        aim_tensor = th.clamp(aim_tensor.detach(), 0, 1)\n",
    "        aim_tensor.requires_grad = True\n",
    "        if loss >= lossn_1:\n",
    "            b += 1\n",
    "            if b > β:\n",
    "                break\n",
    "        else:\n",
    "            b = 0\n",
    "        lossn_1 = loss\n",
    "        if loss < γ:\n",
    "            break\n",
    "    \n",
    "    print(f\"Attack completed at {i} iterations\")\n",
    "    out = model(aim_tensor.detach())\n",
    "    _, pred = th.max(out, 1)\n",
    "    print(pred)\n",
    "    print(f'inverted image {target_label}')\n",
    "    plt.imshow(np.transpose(aim_tensor.detach().reshape(1, 112, 92).numpy() * 255, (1, 2, 0)), cmap=\"Greys\")\n",
    "    plt.show()\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44013f2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's try to reconstruct the data for the first 10 classes (i.e. faces)\n",
    "for cl in range(10):\n",
    "    mi_face(softmax_reg, cl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f14d689",
   "metadata": {},
   "source": [
    "### Exercise: \n",
    "\n",
    "Write the code to try the **model inversion reconstruction** using the `MLP` model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ad83641",
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import MLP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3df4a1b",
   "metadata": {},
   "source": [
    "⚠️ Grab the **pre-trained** weights of the `SoftmaxRegression` model here: [mlp_mia.pt](https://www.dropbox.com/s/8ul2lj2eqcykfxm/mlp_mia.pt?dl=1) and save it into the local `checkpoints` folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58bab294",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp = MLP()\n",
    "weights = load_weights(mlp, model_filename=\"mlp_mia.pt\")\n",
    "if weights is not None:\n",
    "    mlp.load_state_dict(weights)\n",
    "\n",
    "mlp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76662b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reconstruction Attack code HERE\n",
    "for cl in range(10):\n",
    "    mi_face(mlp, cl)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}