docs/source/notebooks/aijack_membership_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yM0rlSybalsu"
   },
   "source": [
    "# Memership Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "AGYSjQHta5tl"
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.svm import SVC\n",
    "from sklearn.metrics import accuracy_score\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "from aijack.attack.membership import ShadowMembershipInferenceAttack\n",
    "from aijack.utils.utils import TorchClassifier, NumpyDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "LEJoOnDjbFC9"
   },
   "outputs": [],
   "source": [
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "\n",
    "df = pd.read_csv(\"/content/sample_data/mnist_train_small.csv\", header=None)\n",
    "\n",
    "X = df[range(1, 785)].values.astype(\"float64\") / 255\n",
    "y = df[0].values\n",
    "\n",
    "# We use the train dataset to train the victim model. The attacker utilize shadow dataset to\n",
    "# prepare membership inference attack. The test dataset is used to evaluate the result of attack.\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X, y, test_size=1 / 3, random_state=42\n",
    ")\n",
    "X_train, X_shadow, y_train, y_shadow = train_test_split(\n",
    "    X_train, y_train, test_size=1 / 2, random_state=42\n",
    ")\n",
    "# We simulate the situation where the distribution of training dataset is different from the test/shadow datasets.\n",
    "X_test = 0.5 * X_test + 0.5 * np.random.normal(size=(X_test.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5vYkm8_sbIA1"
   },
   "source": [
    "## sklearn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zGie8VKUbHBH",
    "outputId": "f972721d-9324-4bad-b30f-ffeb1cf3e067"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.9848484848484849, 0.11054447277636119)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Train the victim\n",
    "\n",
    "clf = SVC(probability=True)\n",
    "clf.fit(X_train, y_train)\n",
    "clf.score(X_train, y_train), clf.score(X_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "2TO-E2CpbUqc"
   },
   "outputs": [],
   "source": [
    "# Train the attacker\n",
    "\n",
    "shadow_models = [SVC(probability=True) for _ in range(2)]\n",
    "attack_models = [SVC(probability=True) for _ in range(10)]\n",
    "\n",
    "attacker = ShadowMembershipInferenceAttack(clf, shadow_models, attack_models)\n",
    "attacker.fit(X_shadow, y_shadow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kfDPXTRebmMO",
    "outputId": "5c4d20f9-9a63-4344-f282-60049ce8a64f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9738243456086402"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the attack result of membership inference\n",
    "in_result = attacker.predict(clf.predict_proba(X_train), y_train)\n",
    "out_result = attacker.predict(clf.predict_proba(X_test), y_test)\n",
    "\n",
    "in_label = np.ones(in_result.shape[0])\n",
    "out_label = np.zeros(out_result.shape[0])\n",
    "\n",
    "accuracy_score(\n",
    "    np.concatenate([in_label, out_label]), np.concatenate([in_result, out_result])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qrDqgKMRbKl1"
   },
   "source": [
    "## PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "hR8Ophn3bLkt"
   },
   "outputs": [],
   "source": [
    "class LM(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LM, self).__init__()\n",
    "        self.lin1 = nn.Linear(28 * 28, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.lin1(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MmG9HrJgceCd",
    "outputId": "8aa8c4b9-1de7-41bb-bfd2-01ddbb8c1066"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.974947494749475, 0.3077846107694615)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Train the victim\n",
    "device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "net = LM().to(torch.double).to(device)\n",
    "optimizer = optim.Adam(net.parameters(), lr=0.001)\n",
    "# You need to wrap the torch module with TorchClassifier\n",
    "clf = TorchClassifier(\n",
    "    net, criterion, optimizer, batch_size=64, epoch=100, device=device\n",
    ")\n",
    "\n",
    "clf.fit(X_train, y_train)\n",
    "clf.score(X_train, y_train), clf.score(X_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "arXgvnJ6cojI"
   },
   "outputs": [],
   "source": [
    "# Train the attacker\n",
    "\n",
    "\n",
    "def create_clf():\n",
    "    _net = LM().to(torch.double).to(device)\n",
    "    _optimizer = optim.Adam(_net.parameters(), lr=0.001)\n",
    "    return TorchClassifier(\n",
    "        _net, criterion, _optimizer, batch_size=64, epoch=100, device=device\n",
    "    )\n",
    "\n",
    "\n",
    "shadow_models = [create_clf() for _ in range(2)]\n",
    "attack_models = [SVC(probability=True) for _ in range(10)]\n",
    "\n",
    "attacker = ShadowMembershipInferenceAttack(clf, shadow_models, attack_models)\n",
    "attacker.fit(X_shadow, y_shadow)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QjIOUDnwc1eC",
    "outputId": "9e7e408a-f98d-4bc3-fb99-08367e02db2f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6735168379209481"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the attack result of membership inference\n",
    "in_result = attacker.predict(clf.predict_proba(X_train), y_train)\n",
    "out_result = attacker.predict(clf.predict_proba(X_test), y_test)\n",
    "\n",
    "in_label = np.ones(in_result.shape[0])\n",
    "out_label = np.zeros(out_result.shape[0])\n",
    "\n",
    "accuracy_score(\n",
    "    np.concatenate([in_label, out_label]), np.concatenate([in_result, out_result])\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}