docs/source/notebooks/aijack_backdoor_federated_learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "109ed59f-0c58-4e27-8e58-e795e868e809",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Backdoor Attack against Federated Learning\n",
    "\n",
    "`Bagdasaryan, Eugene, et al. \"How to backdoor federated learning.\" International conference on artificial intelligence and statistics. PMLR, 2020.`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b1e7d184-ea0b-4e42-a985-faaa936d4264",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2023-11-06 18:58:33--  https://archive.ics.uci.edu/ml/machine-learning-databases/00233/CNAE-9.data\n",
      "Resolving archive.ics.uci.edu (archive.ics.uci.edu)... 128.195.10.252\n",
      "Connecting to archive.ics.uci.edu (archive.ics.uci.edu)|128.195.10.252|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: unspecified\n",
      "Saving to: ‘CNAE-9.data’\n",
      "\n",
      "CNAE-9.data             [            <=>     ]   1.76M   660KB/s    in 2.7s    \n",
      "\n",
      "2023-11-06 18:58:36 (660 KB/s) - ‘CNAE-9.data’ saved [1851120]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://archive.ics.uci.edu/ml/machine-learning-databases/00233/CNAE-9.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1f5c0681-2439-42eb-aaee-dacb2f8e705f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import TensorDataset\n",
    "import random\n",
    "from matplotlib import pyplot as plt\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from aijack.attack.backdoor.modelreplacement import ModelReplacementAttackClientManager\n",
    "from aijack.collaborative.fedavg import FedAVGClient, FedAVGServer, FedAVGAPI\n",
    "\n",
    "batch_size = 1\n",
    "num_rounds = 15\n",
    "lr = 0.0001\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "torch.manual_seed(42)\n",
    "random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ae8053b5-d333-4123-a9ca-f024f5e1f43a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def evaluate_gloal_model(dataloader, client_id=-1):\n",
    "    def _evaluate_global_model(api):\n",
    "        test_loss = 0\n",
    "        correct = 0\n",
    "        with torch.no_grad():\n",
    "            for data, target in dataloader:\n",
    "                data, target = data.to(api.device), target.to(api.device)\n",
    "                if client_id == -1:\n",
    "                    output = api.server(data)\n",
    "                else:\n",
    "                    output = api.clients[client_id](data)\n",
    "                test_loss += F.nll_loss(\n",
    "                    output, target, reduction=\"sum\"\n",
    "                ).item()  # sum up batch loss\n",
    "                pred = output.argmax(\n",
    "                    dim=1, keepdim=True\n",
    "                )  # get the index of the max log-probability\n",
    "                correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "        test_loss /= len(dataloader.dataset)\n",
    "        accuracy = 100.0 * correct / len(dataloader.dataset)\n",
    "        print(f\"Test set: Average loss: {test_loss}, Accuracy: {accuracy}\")\n",
    "\n",
    "        test_loss = 0\n",
    "        correct = 0\n",
    "        with torch.no_grad():\n",
    "            for data, target in dataloader:\n",
    "                data, target = data.to(api.device), target.to(api.device)\n",
    "                data[:, 0] = -1  # inject backdoor\n",
    "                target = torch.zeros_like(target)\n",
    "                if client_id == -1:\n",
    "                    output = api.server(data)\n",
    "                else:\n",
    "                    output = api.clients[client_id](data)\n",
    "                test_loss += F.nll_loss(\n",
    "                    output, target, reduction=\"sum\"\n",
    "                ).item()  # sum up batch loss\n",
    "                pred = output.argmax(\n",
    "                    dim=1, keepdim=True\n",
    "                )  # get the index of the max log-probability\n",
    "                correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "        test_loss /= len(dataloader.dataset)\n",
    "        accuracy = 100.0 * correct / len(dataloader.dataset)\n",
    "        print(f\"Poisoned set: Average loss: {test_loss}, Accuracy: {accuracy}\")\n",
    "\n",
    "    return _evaluate_global_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "03849848-3e86-4d89-b129-6311109f6708",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"CNAE-9.data\", header=None)\n",
    "X = df[range(1, df.shape[1])].values\n",
    "y = df[[0]].values.reshape(-1) - 1\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X, y, test_size=0.2, random_state=42, shuffle=True\n",
    ")\n",
    "\n",
    "X_train = torch.Tensor(X_train)\n",
    "X_test = torch.Tensor(X_test)\n",
    "y_train = torch.Tensor(y_train).to(torch.long)\n",
    "y_test = torch.Tensor(y_test).to(torch.long)\n",
    "\n",
    "poisoned_idx = random.sample(\n",
    "    list(range(int(X_train.shape[0] / 2))), int(0.1 * int(X_train.shape[0] / 2))\n",
    ")\n",
    "X_train[poisoned_idx, 0] = -1  # inject backdoor\n",
    "y_train[poisoned_idx] = 0\n",
    "\n",
    "trainset_1 = TensorDataset(\n",
    "    X_train[: int(X_train.shape[0] / 2)], y_train[: int(X_train.shape[0] / 2)]\n",
    ")\n",
    "trainset_2 = TensorDataset(\n",
    "    X_train[int(X_train.shape[0] / 2) :], y_train[int(X_train.shape[0] / 2) :]\n",
    ")\n",
    "testdataset = TensorDataset(X_test, y_test)\n",
    "\n",
    "trainloader_1 = torch.utils.data.DataLoader(trainset_1, batch_size=batch_size)\n",
    "trainloader_2 = torch.utils.data.DataLoader(trainset_2, batch_size=batch_size)\n",
    "local_dataloaders = [trainloader_1, trainloader_2]\n",
    "test_dataloader = torch.utils.data.DataLoader(testdataset, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8f03124f-fba5-409c-819e-2c9105a1f6ff",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0, 1, 2, 3, 4, 5, 6, 7, 8]),\n",
       " array([132,  99,  89,  86,  84,  94,  90,  96,  94]))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.unique(y_train, return_counts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c0674889-43e6-4854-90f1-371e56786074",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(df.shape[1] - 1, int(df.shape[1] / 2))\n",
    "        self.fc2 = nn.Linear(int(df.shape[1] / 2), 9)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.fc1(x).relu()\n",
    "        out = self.fc2(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(df.shape[1] - 1, 9)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.fc1(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c82c3f1d-76e5-4986-8de1-059745759a12",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "communication 0, epoch 0: client-1 2.1662394603093467\n",
      "communication 0, epoch 0: client-2 2.1922271141299494\n",
      "Test set: Average loss: -5.4690259304587485, Accuracy: 80.55555555555556\n",
      "Poisoned set: Average loss: -3.702957514811445, Accuracy: 38.425925925925924\n",
      "communication 1, epoch 0: client-1 0.6799312920597471\n",
      "communication 1, epoch 0: client-2 0.5005580799708169\n",
      "Test set: Average loss: -5.493952719701661, Accuracy: 89.35185185185185\n",
      "Poisoned set: Average loss: -3.5129815687735877, Accuracy: 33.333333333333336\n",
      "communication 2, epoch 0: client-1 0.5290323983713447\n",
      "communication 2, epoch 0: client-2 0.35746124050794126\n",
      "Test set: Average loss: -5.363429322700809, Accuracy: 70.37037037037037\n",
      "Poisoned set: Average loss: -7.316975637718484, Accuracy: 74.07407407407408\n",
      "communication 3, epoch 0: client-1 0.9123193261237467\n",
      "communication 3, epoch 0: client-2 1.0773419617421665\n",
      "Test set: Average loss: -6.41383598420631, Accuracy: 74.53703703703704\n",
      "Poisoned set: Average loss: -5.338864521295936, Accuracy: 25.925925925925927\n",
      "communication 4, epoch 0: client-1 2.6969062964287533\n",
      "communication 4, epoch 0: client-2 3.1166031406329857\n",
      "Test set: Average loss: -6.398117804416904, Accuracy: 68.98148148148148\n",
      "Poisoned set: Average loss: -11.535380529032814, Accuracy: 80.55555555555556\n",
      "communication 5, epoch 0: client-1 2.36450203070146\n",
      "communication 5, epoch 0: client-2 2.8563940612600134\n",
      "Test set: Average loss: -7.483057366753066, Accuracy: 76.85185185185185\n",
      "Poisoned set: Average loss: -7.282920123250396, Accuracy: 38.888888888888886\n",
      "communication 6, epoch 0: client-1 1.8104334803488789\n",
      "communication 6, epoch 0: client-2 1.665984229675933\n",
      "Test set: Average loss: -7.133859022072068, Accuracy: 79.16666666666667\n",
      "Poisoned set: Average loss: -10.943825021938041, Accuracy: 77.31481481481481\n",
      "communication 7, epoch 0: client-1 1.188665873509647\n",
      "communication 7, epoch 0: client-2 1.589228773664098\n",
      "Test set: Average loss: -7.953649083497347, Accuracy: 85.18518518518519\n",
      "Poisoned set: Average loss: -10.764490551418728, Accuracy: 82.87037037037037\n",
      "communication 8, epoch 0: client-1 0.3532013624386631\n",
      "communication 8, epoch 0: client-2 0.38635327980916545\n",
      "Test set: Average loss: -8.302713400235882, Accuracy: 91.20370370370371\n",
      "Poisoned set: Average loss: -9.670636004871792, Accuracy: 72.68518518518519\n",
      "communication 9, epoch 0: client-1 0.16802960019725577\n",
      "communication 9, epoch 0: client-2 0.12908427009462903\n",
      "Test set: Average loss: -8.271086318901292, Accuracy: 93.05555555555556\n",
      "Poisoned set: Average loss: -10.32074565357632, Accuracy: 79.16666666666667\n",
      "communication 10, epoch 0: client-1 0.11566019391113969\n",
      "communication 10, epoch 0: client-2 0.1096428572484203\n",
      "Test set: Average loss: -8.297636664438027, Accuracy: 94.44444444444444\n",
      "Poisoned set: Average loss: -10.347629803198355, Accuracy: 79.16666666666667\n",
      "communication 11, epoch 0: client-1 0.09780469698989323\n",
      "communication 11, epoch 0: client-2 0.08903112635202053\n",
      "Test set: Average loss: -8.314250176289567, Accuracy: 94.9074074074074\n",
      "Poisoned set: Average loss: -10.441773542651424, Accuracy: 79.62962962962963\n",
      "communication 12, epoch 0: client-1 0.0858488295986393\n",
      "communication 12, epoch 0: client-2 0.08158361000467822\n",
      "Test set: Average loss: -8.33344722242543, Accuracy: 94.9074074074074\n",
      "Poisoned set: Average loss: -10.538744891131365, Accuracy: 80.55555555555556\n",
      "communication 13, epoch 0: client-1 0.07867073030948264\n",
      "communication 13, epoch 0: client-2 0.07543254211632118\n",
      "Test set: Average loss: -8.353414006686458, Accuracy: 94.9074074074074\n",
      "Poisoned set: Average loss: -10.62615246242947, Accuracy: 81.01851851851852\n",
      "communication 14, epoch 0: client-1 0.07302954587543843\n",
      "communication 14, epoch 0: client-2 0.07091031768792234\n",
      "Test set: Average loss: -8.37564472095282, Accuracy: 94.9074074074074\n",
      "Poisoned set: Average loss: -10.70359914832645, Accuracy: 81.48148148148148\n"
     ]
    }
   ],
   "source": [
    "manager = ModelReplacementAttackClientManager(alpha=0.99, gamma=1)\n",
    "ModelReplacementAttackFedAVGClient = manager.attach(FedAVGClient)\n",
    "clients = [\n",
    "    ModelReplacementAttackFedAVGClient(\n",
    "        Net(),\n",
    "        user_id=0,\n",
    "        lr=lr,\n",
    "    ),\n",
    "    FedAVGClient(\n",
    "        Net(),\n",
    "        user_id=1,\n",
    "        lr=lr,\n",
    "    ),\n",
    "]\n",
    "\n",
    "local_optimizers = [optim.SGD(client.parameters(), lr=lr) for client in clients]\n",
    "\n",
    "server = FedAVGServer(clients, Net())\n",
    "\n",
    "api = FedAVGAPI(\n",
    "    server,\n",
    "    clients,\n",
    "    criterion,\n",
    "    local_optimizers,\n",
    "    local_dataloaders,\n",
    "    num_communication=num_rounds,\n",
    "    custom_action=evaluate_gloal_model(test_dataloader),\n",
    ")\n",
    "api.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe13827-6091-457f-b9d2-d65d9cfec063",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3809291d-f7ca-495e-8847-60d42fed1519",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}