docs/source/notebooks/aijack_fedmd.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Raof1tSFHmV"
      },
      "source": [
        "# FedMD: Federated Learning with Model Distillation\n",
        "\n",
        "This tutorial implements FedMD (Federated Learning with Model Distillation), proposed in https://arxiv.org/abs/1910.03581. AIJack supports both single-process and MPI as the backend of FedMD. While FedAVG communicates local gradients to collaboratively train a model without sharing local datasets, malicious servers might be able to recover the training data from the shared gradient (see *Gradient-based Model Inversion Attack against Federated Learning* for the detail). In addition, sending and receiving gradients of the model requires much communication power. To solve these challenges, FedMD communicates not gradients but predicted logits on the global dataset and uses the model-distillation method to share each party's knowledge."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "voM-bUr7FP2r"
      },
      "source": [
        "## Single Process"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "BrIO2C4zOzHz"
      },
      "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 mpi4py import MPI\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "from aijack.collaborative.fedmd import FedMDAPI, FedMDClient, FedMDServer\n",
        "from aijack.utils import NumpyDataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "MX9Cu0wmO69p"
      },
      "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",
        "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\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "fix_seed(seed)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "waDPLc2tO3ZO"
      },
      "outputs": [],
      "source": [
        "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",
        "            NumpyDataset(\n",
        "                x=dataset.data.numpy(),\n",
        "                y=dataset.targets.numpy(),\n",
        "                transform=transform,\n",
        "                return_idx=True,\n",
        "            ),\n",
        "            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(\n",
        "            NumpyDataset(\n",
        "                x=dataset.data.numpy(),\n",
        "                y=dataset.targets.numpy(),\n",
        "                transform=transform,\n",
        "                return_idx=True,\n",
        "            ),\n",
        "            batch_size=test_batch_size,\n",
        "        )\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": "d1aSXndyO9rv"
      },
      "outputs": [],
      "source": [
        "dataloaders = [prepare_dataloader(client_size + 1, c) for c in range(client_size + 1)]\n",
        "public_dataloader = dataloaders[0]\n",
        "local_dataloaders = dataloaders[1:]\n",
        "test_dataloader = prepare_dataloader(client_size, -1, train=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ozQ6dObgO_B-",
        "outputId": "da7a56e0-1497-414c-da92-9d0f9e7f2b3f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 1 (public - pretrain): [1.4732259569076684, 1.509599570077829]\n",
            "acc on validation dataset:  {'clients_score': [0.7988, 0.7907]}\n",
            "epoch 1 (local - pretrain): [0.8319099252216351, 0.8403522926397597]\n",
            "acc on validation dataset:  {'clients_score': [0.8431, 0.8406]}\n",
            "epoch 1, client 0: 248.21629917621613\n",
            "epoch 1, client 1: 269.46992498636246\n",
            "epoch=1 acc on local datasets:  {'clients_score': [0.84605, 0.85175]}\n",
            "epoch=1 acc on public dataset:  {'clients_score': [0.84925, 0.8516]}\n",
            "epoch=1 acc on validation dataset:  {'clients_score': [0.8568, 0.8594]}\n",
            "epoch 2, client 0: 348.2699541449547\n",
            "epoch 2, client 1: 364.1900661587715\n",
            "epoch=2 acc on local datasets:  {'clients_score': [0.8508, 0.85555]}\n",
            "epoch=2 acc on public dataset:  {'clients_score': [0.85395, 0.8567]}\n",
            "epoch=2 acc on validation dataset:  {'clients_score': [0.8598, 0.8641]}\n"
          ]
        }
      ],
      "source": [
        "clients = [\n",
        "    FedMDClient(Net().to(device), public_dataloader, output_dim=10, user_id=c)\n",
        "    for c in range(client_size)\n",
        "]\n",
        "local_optimizers = [optim.SGD(client.parameters(), lr=lr) for client in clients]\n",
        "\n",
        "server = FedMDServer(clients, Net().to(device))\n",
        "\n",
        "api = FedMDAPI(\n",
        "    server,\n",
        "    clients,\n",
        "    public_dataloader,\n",
        "    local_dataloaders,\n",
        "    F.nll_loss,\n",
        "    local_optimizers,\n",
        "    test_dataloader,\n",
        "    num_communication=2,\n",
        ")\n",
        "log = api.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LhrAmm6RFNSH"
      },
      "source": [
        "## MPI\n",
        "\n",
        "You can execute FedMD with MPI-backend via `MPIFedMDClientManager`, `MPIFedMDServerManager`, and `MPIFedMDAPI`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "978eu16pT3Ci",
        "outputId": "a9351915-fb0b-4ed8-c6bf-38fef9aed6c8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Overwriting mpi_fedmd.py\n"
          ]
        }
      ],
      "source": [
        "%%writefile mpi_fedmd.py\n",
        "import random\n",
        "from logging import getLogger\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 mpi4py import MPI\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "from aijack.collaborative.fedmd import FedMDAPI, FedMDClient, FedMDServer\n",
        "from aijack.collaborative.fedmd import MPIFedMDAPI, MPIFedMDClientManager, MPIFedMDServerManager\n",
        "from aijack.utils import NumpyDataset, accuracy_torch_dataloader\n",
        "\n",
        "logger = getLogger(__name__)\n",
        "\n",
        "training_batch_size = 64\n",
        "test_batch_size = 64\n",
        "num_rounds = 2\n",
        "lr = 0.001\n",
        "seed = 0\n",
        "\n",
        "\n",
        "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",
        "            NumpyDataset(x=dataset.data.numpy(), y=dataset.targets.numpy(), transform=transform, return_idx=True),\n",
        "             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(NumpyDataset(x=dataset.data.numpy(), y=dataset.targets.numpy(), transform=transform, return_idx=True),\n",
        "                                                  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\n",
        "\n",
        "def main():\n",
        "    fix_seed(seed)\n",
        "\n",
        "    comm = MPI.COMM_WORLD\n",
        "    myid = comm.Get_rank()\n",
        "    size = comm.Get_size()\n",
        "\n",
        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    model = Net()\n",
        "    model = model.to(device)\n",
        "    optimizer = optim.SGD(model.parameters(), lr=lr)\n",
        "\n",
        "    public_dataloader = prepare_dataloader(size - 1, 0, train=True)\n",
        "\n",
        "    if myid == 0:\n",
        "        dataloader = prepare_dataloader(size + 1, myid+1, train=False)\n",
        "        client_ids = list(range(1, size))\n",
        "        mpi_manager = MPIFedMDServerManager()\n",
        "        MPIFedMDServer = mpi_manager.attach(FedMDServer)\n",
        "        server = MPIFedMDServer(comm, [1, 2], model)\n",
        "        api = MPIFedMDAPI(\n",
        "            comm,\n",
        "            server,\n",
        "            True,\n",
        "            F.nll_loss,\n",
        "            None,\n",
        "            None,\n",
        "            num_communication=num_rounds,\n",
        "            device=device\n",
        "        )\n",
        "    else:\n",
        "        dataloader = prepare_dataloader(size + 1, myid + 1, train=True)\n",
        "        mpi_manager = MPIFedMDClientManager()\n",
        "        MPIFedMDClient = mpi_manager.attach(FedMDClient)\n",
        "        client = MPIFedMDClient(comm, model, public_dataloader, output_dim=10, user_id=myid)\n",
        "        api = MPIFedMDAPI(\n",
        "            comm,\n",
        "            client,\n",
        "            False,\n",
        "            F.nll_loss,\n",
        "            optimizer,\n",
        "            dataloader,\n",
        "            public_dataloader,\n",
        "            num_communication=num_rounds,\n",
        "            device=device\n",
        "        )\n",
        "\n",
        "    api.run()\n",
        "\n",
        "    if myid != 0:\n",
        "      print(f\"client_id={myid}: Accuracy on local dataset is \", accuracy_torch_dataloader(client, dataloader))\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hv4eSXSud8mg",
        "outputId": "fb15bbcc-7d23-4310-bd9c-c6d2417e50ac"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "client_id=2: Accuracy on local dataset is  0.8587333333333333\n",
            "client_id=1: Accuracy on local dataset is  0.8579333333333333\n"
          ]
        }
      ],
      "source": [
        "!sudo mpiexec -np 3 --allow-run-as-root python mpi_fedmd.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8juSW9GjZ5Kx"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "machine_shape": "hm",
      "provenance": []
    },
    "gpuClass": "standard",
    "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 (tags/v3.9.1:1e5d33e, Dec  7 2020, 17:08:21) [MSC v.1927 64 bit (AMD64)]"
    },
    "vscode": {
      "interpreter": {
        "hash": "caa2b01f75ba60e629eaa9e4dabde0c46b243c9a0484934eeb17ad8b3fc9c91a"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}