FL.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "FL.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "yTe3yAbh2V-v"
      },
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "\n",
        "from clients_attackers import *\n",
        "from server import Server\n",
        "from tasks import mnist"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I54AighG3J3z"
      },
      "source": [
        "# FL training configurations\n",
        "AR='fedavg' ## aggregation rule\n",
        "loader_type='dirichlet'\n",
        "loader_path='./data'\n",
        "num_clients=10\n",
        "test_batch_size=64\n",
        "lr=0.01\n",
        "momentum=0.5\n",
        "epochs=30\n",
        "inner_epochs=1\n",
        "device='cuda'\n"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o7bgxPs93LrK"
      },
      "source": [
        "### Specifying the data and model for the Federated Learning task"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6JqbSEfh3RmJ",
        "outputId": "b0e1c68a-8d29-44a3-872a-4b8ba5e617e0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "trainData = mnist.train_dataloader(num_clients, loader_type=loader_type, path=loader_path,\n",
        "                                    store=False)\n",
        "testData = mnist.test_dataloader(test_batch_size)\n",
        "Net = mnist.Net\n",
        "criterion = F.cross_entropy"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Initialize a data loader\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lquZdtVJ5n7G"
      },
      "source": [
        "### Configuring the server"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CpPZqTzk4cu3"
      },
      "source": [
        "model0 = Net()\n",
        "server = Server(model0, testData, criterion, device)\n",
        "server.set_AR(AR)"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oPeeRVuv5sQi"
      },
      "source": [
        "### Configuring the clients and attach them to the server\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Wms8PMC4xgh"
      },
      "source": [
        "list_of_backdoor_attacker=[1,3,5]\n",
        " # i is the client ID\n",
        "for i in range(num_clients):\n",
        "        model = Net()\n",
        "        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)\n",
        "\n",
        "        if i in list_of_backdoor_attacker:           \n",
        "            client_i = Attacker_Backdoor(i, model, trainData[i], optimizer, criterion, device, inner_epochs)\n",
        "        else:\n",
        "            client_i = Client(i, model, trainData[i], optimizer, criterion, device, inner_epochs)\n",
        "        server.attach(client_i)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vhhKIROI6Aha"
      },
      "source": [
        "### Sanity check of the pipeline"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MB3KEUs_6EQq",
        "outputId": "ecc4d085-30f7-4c59-d798-98ca74b51dc2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "loss, accuracy = server.test()\n",
        "aloss, asr     = server.test_backdoor()\n",
        "steps=0\n",
        "print(f\"[{steps}]\\t  Loss:\\t{ loss:.4f}\\t\\tAccuracy:\\t{accuracy:.2f}\") \n",
        "print(f\"[{steps}]\\tA.Loss:\\t{aloss:.4f}\\t\\t     ASR:\\t{asr:.2f}\\n\\n\") "
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 2.3059, Accuracy: 560/10000 (6%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.2624, Success rate: 79/10000 (1%)\n",
            "\n",
            "[0]\t  Loss:\t2.3059\t\tAccuracy:\t5.60\n",
            "[0]\tA.Loss:\t2.2624\t\t     ASR:\t0.79\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "75_-6Xha7Mme"
      },
      "source": [
        "### Run the training for the Federated learning task"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sRUW-nRc6NaN",
        "outputId": "8fc58003-84d1-4a58-8740-ddf814e56cfd",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "for j in range(epochs):\n",
        "    steps = j + 1\n",
        "    server.distribute()\n",
        "    #         group=Random().sample(range(5),1) # random subset sampling was not used in our paper\n",
        "    group = range(num_clients) # we assume all clients are involved in synchronization\n",
        "    server.train(group)\n",
        "\n",
        "    loss, accuracy = server.test()\n",
        "    aloss, asr     = server.test_backdoor()\n",
        "\n",
        "    print(f\"[{steps}]\\t  Loss:\\t{ loss:.4f}\\t\\tAccuracy:\\t{accuracy:.2f}\") \n",
        "    print(f\"[{steps}]\\tA.Loss:\\t{aloss:.4f}\\t\\t     ASR:\\t{asr:.2f}\\n\\n\") "
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 2.2952, Accuracy: 1032/10000 (10%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.0398, Success rate: 10000/10000 (100%)\n",
            "\n",
            "[1]\t  Loss:\t2.2952\t\tAccuracy:\t10.32\n",
            "[1]\tA.Loss:\t2.0398\t\t     ASR:\t100.00\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 2.2397, Accuracy: 1032/10000 (10%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 1.5253, Success rate: 10000/10000 (100%)\n",
            "\n",
            "[2]\t  Loss:\t2.2397\t\tAccuracy:\t10.32\n",
            "[2]\tA.Loss:\t1.5253\t\t     ASR:\t100.00\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 1.7479, Accuracy: 3029/10000 (30%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 1.4667, Success rate: 7233/10000 (72%)\n",
            "\n",
            "[3]\t  Loss:\t1.7479\t\tAccuracy:\t30.29\n",
            "[3]\tA.Loss:\t1.4667\t\t     ASR:\t72.33\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.9148, Accuracy: 7455/10000 (75%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.2911, Success rate: 2214/10000 (22%)\n",
            "\n",
            "[4]\t  Loss:\t0.9148\t\tAccuracy:\t74.55\n",
            "[4]\tA.Loss:\t2.2911\t\t     ASR:\t22.14\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.6428, Accuracy: 8164/10000 (82%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.6497, Success rate: 1807/10000 (18%)\n",
            "\n",
            "[5]\t  Loss:\t0.6428\t\tAccuracy:\t81.64\n",
            "[5]\tA.Loss:\t2.6497\t\t     ASR:\t18.07\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.5216, Accuracy: 8508/10000 (85%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.9475, Success rate: 1566/10000 (16%)\n",
            "\n",
            "[6]\t  Loss:\t0.5216\t\tAccuracy:\t85.08\n",
            "[6]\tA.Loss:\t2.9475\t\t     ASR:\t15.66\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.4354, Accuracy: 8784/10000 (88%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 3.2397, Success rate: 1310/10000 (13%)\n",
            "\n",
            "[7]\t  Loss:\t0.4354\t\tAccuracy:\t87.84\n",
            "[7]\tA.Loss:\t3.2397\t\t     ASR:\t13.10\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.4277, Accuracy: 8807/10000 (88%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.8063, Success rate: 1555/10000 (16%)\n",
            "\n",
            "[8]\t  Loss:\t0.4277\t\tAccuracy:\t88.07\n",
            "[8]\tA.Loss:\t2.8063\t\t     ASR:\t15.55\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.3808, Accuracy: 8945/10000 (89%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 3.0221, Success rate: 1356/10000 (14%)\n",
            "\n",
            "[9]\t  Loss:\t0.3808\t\tAccuracy:\t89.45\n",
            "[9]\tA.Loss:\t3.0221\t\t     ASR:\t13.56\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.3566, Accuracy: 8983/10000 (90%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 3.1053, Success rate: 1330/10000 (13%)\n",
            "\n",
            "[10]\t  Loss:\t0.3566\t\tAccuracy:\t89.83\n",
            "[10]\tA.Loss:\t3.1053\t\t     ASR:\t13.30\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.3652, Accuracy: 8957/10000 (90%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.7327, Success rate: 1523/10000 (15%)\n",
            "\n",
            "[11]\t  Loss:\t0.3652\t\tAccuracy:\t89.57\n",
            "[11]\tA.Loss:\t2.7327\t\t     ASR:\t15.23\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.3269, Accuracy: 9085/10000 (91%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 3.0351, Success rate: 1345/10000 (13%)\n",
            "\n",
            "[12]\t  Loss:\t0.3269\t\tAccuracy:\t90.85\n",
            "[12]\tA.Loss:\t3.0351\t\t     ASR:\t13.45\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.3068, Accuracy: 9147/10000 (91%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 3.0451, Success rate: 1301/10000 (13%)\n",
            "\n",
            "[13]\t  Loss:\t0.3068\t\tAccuracy:\t91.47\n",
            "[13]\tA.Loss:\t3.0451\t\t     ASR:\t13.01\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.3129, Accuracy: 9107/10000 (91%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.7852, Success rate: 1422/10000 (14%)\n",
            "\n",
            "[14]\t  Loss:\t0.3129\t\tAccuracy:\t91.07\n",
            "[14]\tA.Loss:\t2.7852\t\t     ASR:\t14.22\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2832, Accuracy: 9193/10000 (92%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 3.0340, Success rate: 1300/10000 (13%)\n",
            "\n",
            "[15]\t  Loss:\t0.2832\t\tAccuracy:\t91.93\n",
            "[15]\tA.Loss:\t3.0340\t\t     ASR:\t13.00\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2887, Accuracy: 9207/10000 (92%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.7805, Success rate: 1420/10000 (14%)\n",
            "\n",
            "[16]\t  Loss:\t0.2887\t\tAccuracy:\t92.07\n",
            "[16]\tA.Loss:\t2.7805\t\t     ASR:\t14.20\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2623, Accuracy: 9277/10000 (93%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 3.0507, Success rate: 1283/10000 (13%)\n",
            "\n",
            "[17]\t  Loss:\t0.2623\t\tAccuracy:\t92.77\n",
            "[17]\tA.Loss:\t3.0507\t\t     ASR:\t12.83\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2602, Accuracy: 9290/10000 (93%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.8132, Success rate: 1391/10000 (14%)\n",
            "\n",
            "[18]\t  Loss:\t0.2602\t\tAccuracy:\t92.90\n",
            "[18]\tA.Loss:\t2.8132\t\t     ASR:\t13.91\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2470, Accuracy: 9315/10000 (93%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.8774, Success rate: 1349/10000 (13%)\n",
            "\n",
            "[19]\t  Loss:\t0.2470\t\tAccuracy:\t93.15\n",
            "[19]\tA.Loss:\t2.8774\t\t     ASR:\t13.49\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2327, Accuracy: 9356/10000 (94%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.9959, Success rate: 1285/10000 (13%)\n",
            "\n",
            "[20]\t  Loss:\t0.2327\t\tAccuracy:\t93.56\n",
            "[20]\tA.Loss:\t2.9959\t\t     ASR:\t12.85\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2443, Accuracy: 9323/10000 (93%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.6173, Success rate: 1512/10000 (15%)\n",
            "\n",
            "[21]\t  Loss:\t0.2443\t\tAccuracy:\t93.23\n",
            "[21]\tA.Loss:\t2.6173\t\t     ASR:\t15.12\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2297, Accuracy: 9361/10000 (94%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.6718, Success rate: 1467/10000 (15%)\n",
            "\n",
            "[22]\t  Loss:\t0.2297\t\tAccuracy:\t93.61\n",
            "[22]\tA.Loss:\t2.6718\t\t     ASR:\t14.67\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2165, Accuracy: 9417/10000 (94%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.7153, Success rate: 1423/10000 (14%)\n",
            "\n",
            "[23]\t  Loss:\t0.2165\t\tAccuracy:\t94.17\n",
            "[23]\tA.Loss:\t2.7153\t\t     ASR:\t14.23\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2009, Accuracy: 9451/10000 (95%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.8613, Success rate: 1374/10000 (14%)\n",
            "\n",
            "[24]\t  Loss:\t0.2009\t\tAccuracy:\t94.51\n",
            "[24]\tA.Loss:\t2.8613\t\t     ASR:\t13.74\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.1925, Accuracy: 9475/10000 (95%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.7555, Success rate: 1385/10000 (14%)\n",
            "\n",
            "[25]\t  Loss:\t0.1925\t\tAccuracy:\t94.75\n",
            "[25]\tA.Loss:\t2.7555\t\t     ASR:\t13.85\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.2007, Accuracy: 9463/10000 (95%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.2352, Success rate: 1714/10000 (17%)\n",
            "\n",
            "[26]\t  Loss:\t0.2007\t\tAccuracy:\t94.63\n",
            "[26]\tA.Loss:\t2.2352\t\t     ASR:\t17.14\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.1847, Accuracy: 9513/10000 (95%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.3875, Success rate: 1582/10000 (16%)\n",
            "\n",
            "[27]\t  Loss:\t0.1847\t\tAccuracy:\t95.13\n",
            "[27]\tA.Loss:\t2.3875\t\t     ASR:\t15.82\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.1798, Accuracy: 9515/10000 (95%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 2.1332, Success rate: 1913/10000 (19%)\n",
            "\n",
            "[28]\t  Loss:\t0.1798\t\tAccuracy:\t95.15\n",
            "[28]\tA.Loss:\t2.1332\t\t     ASR:\t19.13\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.1723, Accuracy: 9533/10000 (95%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 1.9405, Success rate: 2262/10000 (23%)\n",
            "\n",
            "[29]\t  Loss:\t0.1723\t\tAccuracy:\t95.33\n",
            "[29]\tA.Loss:\t1.9405\t\t     ASR:\t22.62\n",
            "\n",
            "\n",
            "[Server] Start testing\n",
            "\n",
            "Test set: Average loss: 0.1658, Accuracy: 9561/10000 (96%)\n",
            "\n",
            "[Server] Start testing backdoor\n",
            "\n",
            "Test set (Backdoored): Average loss: 1.7255, Success rate: 2651/10000 (27%)\n",
            "\n",
            "[30]\t  Loss:\t0.1658\t\tAccuracy:\t95.61\n",
            "[30]\tA.Loss:\t1.7255\t\t     ASR:\t26.51\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}