docs/source/notebooks/aijack_poisoning_federated_learning.ipynb

{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "SROAZ9dO80s8"
            },
            "source": [
                "# Poisoning Attack against Federated Learning\n",
                "\n",
                "This tutorial demonstrates that malicious clients effectively decrease the performance of the final global model by injecting noise into the local update or data."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {
                "id": "J4s0w9rHwOd8"
            },
            "outputs": [],
            "source": [
                "import random\n",
                "\n",
                "import numpy as np\n",
                "import torch\n",
                "import torch.nn as nn\n",
                "import torch.nn.functional as F\n",
                "import torch.optim as optim\n",
                "from torchvision import datasets, transforms\n",
                "\n",
                "from aijack.attack.poison import HistoryAttackClientWrapper\n",
                "from aijack.attack.poison import LabelFlipAttackClientManager\n",
                "from aijack.attack.poison import MAPFClientWrapper\n",
                "from aijack.collaborative.fedavg import FedAVGClient, FedAVGServer, FedAVGAPI\n",
                "\n",
                "\n",
                "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",
                "    return _evaluate_global_model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {
                "id": "idlkuHjjULVz"
            },
            "outputs": [],
            "source": [
                "training_batch_size = 64\n",
                "test_batch_size = 64\n",
                "num_rounds = 5\n",
                "lr = 0.001\n",
                "seed = 0\n",
                "client_size = 2\n",
                "criterion = F.nll_loss"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {
                "id": "gsh_CjamwqfV"
            },
            "outputs": [],
            "source": [
                "def fix_seed(seed):\n",
                "    random.seed(seed)\n",
                "    np.random.seed(seed)\n",
                "    torch.manual_seed(seed)\n",
                "    torch.cuda.manual_seed_all(seed)\n",
                "    torch.backends.cudnn.deterministic = True\n",
                "\n",
                "\n",
                "def prepare_dataloader(num_clients, myid, train=True, path=\"\"):\n",
                "    transform = transforms.Compose(\n",
                "        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
                "    )\n",
                "    if train:\n",
                "        dataset = datasets.MNIST(path, train=True, download=True, transform=transform)\n",
                "        idxs = list(range(len(dataset.data)))\n",
                "        random.shuffle(idxs)\n",
                "        idx = np.array_split(idxs, num_clients, 0)[myid - 1]\n",
                "        dataset.data = dataset.data[idx]\n",
                "        dataset.targets = dataset.targets[idx]\n",
                "        train_loader = torch.utils.data.DataLoader(\n",
                "            dataset, batch_size=training_batch_size\n",
                "        )\n",
                "        return train_loader\n",
                "    else:\n",
                "        dataset = datasets.MNIST(path, train=False, download=True, transform=transform)\n",
                "        test_loader = torch.utils.data.DataLoader(dataset, batch_size=test_batch_size)\n",
                "        return test_loader\n",
                "\n",
                "\n",
                "class Net(nn.Module):\n",
                "    def __init__(self):\n",
                "        super(Net, self).__init__()\n",
                "        self.ln = nn.Linear(28 * 28, 10)\n",
                "\n",
                "    def forward(self, x):\n",
                "        x = self.ln(x.reshape(-1, 28 * 28))\n",
                "        output = F.log_softmax(x, dim=1)\n",
                "        return output"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {
                "id": "UX-5FPlbxs4z"
            },
            "outputs": [],
            "source": [
                "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
                "fix_seed(seed)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 423,
                    "referenced_widgets": [
                        "454b6a9d85304fe8868eb44bc2af1cb6",
                        "e2ffa42cb4d64d78ba654dfe6f4be0d0",
                        "13804a3c3a804729bd83ab0271aaa2f3",
                        "2fe2f1241c614b3bba7b792fd10061a9",
                        "da103b41d41a439b909d8bae77e33e49",
                        "145a7a75341b43f0ab205e45f726f406",
                        "5e8368f05c244084a1505ba533f84c01",
                        "f9efb089952e479e8d3ac121614dbd04",
                        "9b03839134814afa963f2fba0ea0b79e",
                        "4f46ebfdcea1453581fe5814ee464bad",
                        "837bba7ab6ad465d9bd00ca1e1df83cd",
                        "4962b31a3403404ba6bccb43aa9a683d",
                        "7216bafce927421e99a0dcb29e73826c",
                        "6536b22da14a451d824a9ad6b66be104",
                        "da8b4a73bfa44053bf6c64daf37a3340",
                        "42f57e297e4f419fa7733f2fba334305",
                        "29cd1a72188f4547a03f344d8163df05",
                        "fb1aa5ba266b471399aae74cfc670c97",
                        "614c7f033c504ab286fdf0979e5b5d68",
                        "ec37c87a0d9c439a88788285db89dac5",
                        "18b2594fbbfe4304b5cab0169da04e86",
                        "e49e0e37e380480cb17f00792598151b",
                        "ed9f9da03f8347cfa14341c46249f1dd",
                        "00384a239d084e4daccd0b5c345d9da7",
                        "2645063fb9084a2ebc584440e984b841",
                        "bb4769236d7849a3ad17dca8dc6720f7",
                        "64fa697a487b459797f4cc776afa24af",
                        "4495dc9bbe5946a6a93cf1ac7fca4e5c",
                        "b6ec24062be448eab2b1dbac757515c2",
                        "d1868e3a654b43e292e98c9b6475b2bc",
                        "ad309e1030744202a969ad295a7dee18",
                        "584b693b713b436096b347c67062c66d",
                        "4f5b89295ddd41e38f3b3b16e0c8eec0",
                        "7ff6082e6f094439ad0d99479cf5fdb3",
                        "4774b9fc69d94e1b8efee595b3f36016",
                        "98a0f8c74d374edc82d8b5f27151065d",
                        "3017d721c21b4cdd81acff0b27178934",
                        "480690b7e2bd44f0a8cda5f9d6d343c5",
                        "93e587269dc1466e854a16bdc365069e",
                        "aa27b1502f114c6db46e763f5fc181ac",
                        "063a333c198b441ea13563ef984fcd90",
                        "e1720da2edf84fbfaf5ec1ea391f345c",
                        "264419bb7e8f47c987601c71af84adaf",
                        "7f0ceb349cb842059ed6f590a808f798"
                    ]
                },
                "id": "zz_YjoioAcLD",
                "outputId": "8d049b38-e708-4198-d9ce-bf189dbf4ef6"
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
                        "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to MNIST/raw/train-images-idx3-ubyte.gz\n"
                    ]
                },
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "454b6a9d85304fe8868eb44bc2af1cb6",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "  0%|          | 0/9912422 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Extracting MNIST/raw/train-images-idx3-ubyte.gz to MNIST/raw\n",
                        "\n",
                        "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
                        "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to MNIST/raw/train-labels-idx1-ubyte.gz\n"
                    ]
                },
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "4962b31a3403404ba6bccb43aa9a683d",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "  0%|          | 0/28881 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Extracting MNIST/raw/train-labels-idx1-ubyte.gz to MNIST/raw\n",
                        "\n",
                        "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
                        "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to MNIST/raw/t10k-images-idx3-ubyte.gz\n"
                    ]
                },
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "ed9f9da03f8347cfa14341c46249f1dd",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "  0%|          | 0/1648877 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Extracting MNIST/raw/t10k-images-idx3-ubyte.gz to MNIST/raw\n",
                        "\n",
                        "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
                        "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
                    ]
                },
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "7ff6082e6f094439ad0d99479cf5fdb3",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "  0%|          | 0/4542 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Extracting MNIST/raw/t10k-labels-idx1-ubyte.gz to MNIST/raw\n",
                        "\n"
                    ]
                }
            ],
            "source": [
                "local_dataloaders = [prepare_dataloader(client_size, c) for c in range(client_size)]\n",
                "test_dataloader = prepare_dataloader(client_size, -1, train=False)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "hFEcEGO6htgz"
            },
            "source": [
                "## History Attack"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "id": "UEzYT3tex0Nm",
                "outputId": "cce61266-ce62-4190-ffef-26763080705a"
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "communication 0, epoch 0: client-1 0.019623182545105616\n",
                        "communication 0, epoch 0: client-2 0.019723439192771912\n",
                        "Test set: Average loss: 6.538579542136137, Accuracy: 82.57\n",
                        "communication 1, epoch 0: client-1 0.1011283678372701\n",
                        "communication 1, epoch 0: client-2 0.10503993360201518\n",
                        "Test set: Average loss: 109.3780958984375, Accuracy: 24.71\n",
                        "communication 2, epoch 0: client-1 1.307175351079305\n",
                        "communication 2, epoch 0: client-2 1.322490109125773\n",
                        "Test set: Average loss: 514.1672515625, Accuracy: 59.83\n",
                        "communication 3, epoch 0: client-1 7.669983786519368\n",
                        "communication 3, epoch 0: client-2 7.649992772420247\n",
                        "Test set: Average loss: 440.93763037109375, Accuracy: 46.56\n",
                        "communication 4, epoch 0: client-1 6.582273025512695\n",
                        "communication 4, epoch 0: client-2 6.532691622924805\n",
                        "Test set: Average loss: 616.4571529296875, Accuracy: 40.87\n"
                    ]
                }
            ],
            "source": [
                "manager = HistoryAttackClientWrapper(lam=3)\n",
                "HistoryAttackFedAVGClient = manager.attach(FedAVGClient)\n",
                "clients = [\n",
                "    HistoryAttackFedAVGClient(\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().to(device))\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": "markdown",
            "metadata": {
                "id": "cGaWpZSshrf5"
            },
            "source": [
                "## Label Flip Attack"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "id": "GmGoNF8HgpEY",
                "outputId": "1fcf6c5c-1dce-4f9c-c69c-6c6a9bb57cda"
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "communication 0, epoch 0: client-1 0.020543035099903743\n",
                        "communication 0, epoch 0: client-2 0.020125101908047994\n",
                        "Test set: Average loss: 28.122399871826172, Accuracy: 73.83\n",
                        "communication 1, epoch 0: client-1 0.5485365001996358\n",
                        "communication 1, epoch 0: client-2 0.4188099824587504\n",
                        "Test set: Average loss: 364.91654228515625, Accuracy: 39.13\n",
                        "communication 2, epoch 0: client-1 5.370502290852865\n",
                        "communication 2, epoch 0: client-2 5.275297354125977\n",
                        "Test set: Average loss: 1108.481857421875, Accuracy: 35.73\n",
                        "communication 3, epoch 0: client-1 15.66210668334961\n",
                        "communication 3, epoch 0: client-2 16.781931443277994\n",
                        "Test set: Average loss: 1227.238296875, Accuracy: 33.2\n",
                        "communication 4, epoch 0: client-1 16.553591099039714\n",
                        "communication 4, epoch 0: client-2 18.498205289713542\n",
                        "Test set: Average loss: 1096.7404470703125, Accuracy: 42.83\n"
                    ]
                }
            ],
            "source": [
                "manager = LabelFlipAttackClientManager(victim_label=0, target_label=1)\n",
                "LabelFlipAttackFedAVGClient = manager.attach(FedAVGClient)\n",
                "clients = [\n",
                "    LabelFlipAttackFedAVGClient(\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().to(device))\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": "markdown",
            "metadata": {
                "id": "Ogkc6djGhpsT"
            },
            "source": [
                "## MAPF"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "metadata": {
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "id": "JE9boSiRoKWq",
                "outputId": "9abc6bf7-4079-48bf-ab40-044a37305d51"
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "communication 0, epoch 0: client-1 0.019650927847623824\n",
                        "communication 0, epoch 0: client-2 0.019755615478754044\n",
                        "Test set: Average loss: 6.351612661059061, Accuracy: 83.04\n",
                        "communication 1, epoch 0: client-1 0.10415176281332969\n",
                        "communication 1, epoch 0: client-2 0.10801099200248718\n",
                        "Test set: Average loss: 64.09054548339844, Accuracy: 35.71\n",
                        "communication 2, epoch 0: client-1 0.6900041089375814\n",
                        "communication 2, epoch 0: client-2 0.6906570717493693\n",
                        "Test set: Average loss: 423.16109165039063, Accuracy: 50.94\n",
                        "communication 3, epoch 0: client-1 6.43880789159139\n",
                        "communication 3, epoch 0: client-2 6.371651240030925\n",
                        "Test set: Average loss: 596.9880190429687, Accuracy: 41.92\n",
                        "communication 4, epoch 0: client-1 8.962839545694987\n",
                        "communication 4, epoch 0: client-2 8.890711385091146\n",
                        "Test set: Average loss: 597.5907002929688, Accuracy: 39.4\n"
                    ]
                }
            ],
            "source": [
                "manager = MAPFClientWrapper(lam=3)\n",
                "MAPFFedAVGClient = manager.attach(FedAVGClient)\n",
                "clients = [\n",
                "    MAPFFedAVGClient(\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().to(device))\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": 9,
            "metadata": {
                "id": "qA8_TCaHhV6L"
            },
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "colab": {
            "provenance": []
        },
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "name": "python",
            "version": "3.9.1 (tags/v3.9.1:1e5d33e, Dec  7 2020, 17:08:21) [MSC v.1927 64 bit (AMD64)]"
        },
        "vscode": {
            "interpreter": {
                "hash": "caa2b01f75ba60e629eaa9e4dabde0c46b243c9a0484934eeb17ad8b3fc9c91a"
            }
        },
        "widgets": {
            "application/vnd.jupyter.widget-state+json": {}
        }
    },
    "nbformat": 4,
    "nbformat_minor": 0
}