docs/aijack_mid.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "07c070e9",
   "metadata": {},
   "source": [
    "# Mutual Information-based Defense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "loose-negotiation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torchvision.datasets import MNIST\n",
    "from torch.utils.data import TensorDataset, Dataset, DataLoader\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "import numpy as np\n",
    "\n",
    "from aijack.defense import VIB, KL_between_normals, mib_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "surgical-summit",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim_z = 256\n",
    "beta = 1e-3\n",
    "batch_size = 100\n",
    "samples_amount = 15\n",
    "num_epochs = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "developmental-climate",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = MNIST(\"MNIST/.\", download=True, train=True)\n",
    "train_dataset = TensorDataset(\n",
    "    train_data.train_data.view(-1, 28 * 28).float() / 255, train_data.train_labels\n",
    ")\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size)\n",
    "\n",
    "test_data = MNIST(\"MNIST/.\", download=True, train=False)\n",
    "test_dataset = TensorDataset(\n",
    "    test_data.test_data.view(-1, 28 * 28).float() / 255, test_data.test_labels\n",
    ")\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "direct-commission",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = nn.Sequential(\n",
    "    nn.Linear(in_features=784, out_features=1024),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(in_features=1024, out_features=1024),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(in_features=1024, out_features=2 * dim_z),\n",
    ")\n",
    "decoder = nn.Linear(in_features=dim_z, out_features=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "official-syndication",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = VIB(encoder, decoder, dim_z, num_samples=samples_amount)\n",
    "opt = torch.optim.Adam(net.parameters(), lr=1e-4)\n",
    "scheduler = torch.optim.lr_scheduler.ExponentialLR(opt, gamma=0.97)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "automatic-europe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    loss_by_epoch = []\n",
    "    accuracy_by_epoch = []\n",
    "    I_ZX_bound_by_epoch = []\n",
    "    I_ZY_bound_by_epoch = []\n",
    "\n",
    "    loss_by_epoch_test = []\n",
    "    accuracy_by_epoch_test = []\n",
    "    I_ZX_bound_by_epoch_test = []\n",
    "    I_ZY_bound_by_epoch_test = []\n",
    "\n",
    "    if epoch % 2 == 0 and epoch > 0:\n",
    "        scheduler.step()\n",
    "\n",
    "    for x_batch, y_batch in tqdm(train_loader):\n",
    "        x_batch = x_batch\n",
    "        y_batch = y_batch\n",
    "\n",
    "        y_pred, result_dict = net(x_batch)\n",
    "        sampled_y_pred = result_dict[\"sampled_decoded_outputs\"]\n",
    "        p_z_given_x_mu = result_dict[\"p_z_given_x_mu\"]\n",
    "        p_z_given_x_sigma = result_dict[\"p_z_given_x_sigma\"]\n",
    "\n",
    "        approximated_z_mean = torch.zeros_like(p_z_given_x_mu)\n",
    "        approximated_z_sigma = torch.ones_like(p_z_given_x_sigma)\n",
    "\n",
    "        loss, I_ZY_bound, I_ZX_bound = mib_loss(\n",
    "            y_batch,\n",
    "            sampled_y_pred,\n",
    "            p_z_given_x_mu,\n",
    "            p_z_given_x_sigma,\n",
    "            approximated_z_mean,\n",
    "            approximated_z_sigma,\n",
    "            beta=beta,\n",
    "        )\n",
    "\n",
    "        prediction = torch.max(y_pred, dim=1)[1]\n",
    "        accuracy = torch.mean((prediction == y_batch).float())\n",
    "\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        opt.zero_grad()\n",
    "\n",
    "        I_ZX_bound_by_epoch.append(I_ZX_bound.item())\n",
    "        I_ZY_bound_by_epoch.append(I_ZY_bound.item())\n",
    "\n",
    "        loss_by_epoch.append(loss.item())\n",
    "        accuracy_by_epoch.append(accuracy.item())\n",
    "\n",
    "    for x_batch, y_batch in tqdm(test_loader):\n",
    "        x_batch = x_batch\n",
    "        y_batch = y_batch\n",
    "\n",
    "        y_pred, result_dict = net(x_batch)\n",
    "        sampled_y_pred = result_dict[\"sampled_decoded_outputs\"]\n",
    "        p_z_given_x_mu = result_dict[\"p_z_given_x_mu\"]\n",
    "        p_z_given_x_sigma = result_dict[\"p_z_given_x_sigma\"]\n",
    "\n",
    "        approximated_z_mean = torch.zeros_like(p_z_given_x_mu)\n",
    "        approximated_z_sigma = torch.ones_like(p_z_given_x_sigma)\n",
    "\n",
    "        loss, I_ZY_bound, I_ZX_bound = mib_loss(\n",
    "            y_batch,\n",
    "            sampled_y_pred,\n",
    "            p_z_given_x_mu,\n",
    "            p_z_given_x_sigma,\n",
    "            approximated_z_mean,\n",
    "            approximated_z_sigma,\n",
    "            beta=beta,\n",
    "        )\n",
    "\n",
    "        prediction = torch.max(y_pred, dim=1)[1]\n",
    "        accuracy = torch.mean((prediction == y_batch).float())\n",
    "\n",
    "        I_ZX_bound_by_epoch_test.append(I_ZX_bound.item())\n",
    "        I_ZY_bound_by_epoch_test.append(I_ZY_bound.item())\n",
    "\n",
    "        loss_by_epoch_test.append(loss.item())\n",
    "        accuracy_by_epoch_test.append(accuracy.item())\n",
    "\n",
    "    print(\n",
    "        \"epoch\",\n",
    "        epoch,\n",
    "        \"loss\",\n",
    "        np.mean(loss_by_epoch_test),\n",
    "        \"prediction\",\n",
    "        np.mean(accuracy_by_epoch_test),\n",
    "    )\n",
    "\n",
    "    print(\n",
    "        \"I_ZX_bound\",\n",
    "        np.mean(I_ZX_bound_by_epoch_test),\n",
    "        \"I_ZY_bound\",\n",
    "        np.mean(I_ZY_bound_by_epoch_test),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "intensive-gnome",
   "metadata": {},
   "outputs": [],
   "source": [
    "from aijack.attack import GradientInversion_Attack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "annoying-auction",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred, result_dict = net(x_batch[:1])\n",
    "sampled_y_pred = result_dict[\"sampled_decoded_outputs\"]\n",
    "p_z_given_x_mu = result_dict[\"p_z_given_x_mu\"]\n",
    "p_z_given_x_sigma = result_dict[\"p_z_given_x_sigma\"]\n",
    "\n",
    "approximated_z_mean = torch.zeros_like(p_z_given_x_mu)\n",
    "approximated_z_sigma = torch.ones_like(p_z_given_x_sigma)\n",
    "\n",
    "\n",
    "loss, I_ZY_bound, I_ZX_bound = mib_loss(\n",
    "    y_batch[:1],\n",
    "    sampled_y_pred,\n",
    "    p_z_given_x_mu,\n",
    "    p_z_given_x_sigma,\n",
    "    approximated_z_mean,\n",
    "    approximated_z_sigma,\n",
    "    beta=beta,\n",
    ")\n",
    "\n",
    "received_gradients = torch.autograd.grad(loss, net.parameters())\n",
    "received_gradients = [cg.detach() for cg in received_gradients]\n",
    "received_gradients = [cg for cg in received_gradients]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "national-treasurer",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import cv2\n",
    "\n",
    "net.eval()\n",
    "cpl_attacker = GradientInversion_Attack(\n",
    "    net,\n",
    "    (784,),\n",
    "    lr=0.3,\n",
    "    log_interval=50,\n",
    "    optimizer_class=torch.optim.LBFGS,\n",
    "    distancename=\"l2\",\n",
    "    optimize_label=False,\n",
    "    num_iteration=200,\n",
    ")\n",
    "\n",
    "num_seeds = 5\n",
    "fig = plt.figure(figsize=(6, 2))\n",
    "for s in tqdm(range(num_seeds)):\n",
    "    cpl_attacker.reset_seed(s)\n",
    "    try:\n",
    "        result = cpl_attacker.attack(received_gradients)\n",
    "        ax1 = fig.add_subplot(2, num_seeds, s + 1)\n",
    "        ax1.imshow(result[0].cpu().detach().numpy()[0].reshape(28, 28), cmap=\"gray\")\n",
    "        ax1.axis(\"off\")\n",
    "        ax1.set_title(torch.argmax(result[1]).cpu().item())\n",
    "        ax2 = fig.add_subplot(2, num_seeds, num_seeds + s + 1)\n",
    "        ax2.imshow(\n",
    "            cv2.medianBlur(result[0].cpu().detach().numpy()[0].reshape(28, 28), 5),\n",
    "            cmap=\"gray\",\n",
    "        )\n",
    "        ax2.axis(\"off\")\n",
    "    except:\n",
    "        pass\n",
    "plt.suptitle(\"Result of CPL\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}