5_9_9_FederatedLearning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "5-9-9-FederatedLearning.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pujOZk0HbuDS"
      },
      "source": [
        "# Federated Learning Walkthrough\n",
        "\n",
        "In this notebook you will get hands on experience programming a federated learning simulation. Note that, this notebook, through programming abstractions, simulates the clients and servers that operate in a federated learning setup. This notebook serves as an example to illustrate the steps taken in a federated learning system.\n",
        "\n",
        "First, let's import our libraries, load the mnist dataset and define our model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "do7yY5OndbRk"
      },
      "source": [
        "# Import modules\n",
        "import os\n",
        "import sys\n",
        "import torch\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.datasets as datasets\n",
        "import torch.nn as nn\n",
        "from torch.autograd import Variable\n",
        "import torch.optim as optim\n",
        "import random\n",
        "import torch.nn.functional as F\n",
        "import copy\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sleBEUqLcoZr"
      },
      "source": [
        "# Load MNIST Dataset\n",
        "d = './data'\n",
        "if not os.path.exists(d):\n",
        "    os.mkdir(d)\n",
        "    \n",
        "trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])\n",
        "train_set = datasets.MNIST(root=d, train=True, transform=trans, download=True)\n",
        "test_set = datasets.MNIST(root=d, train=False, transform=trans, download=True)\n",
        "\n",
        "batch_size = 32\n",
        "global_train_loader = torch.utils.data.DataLoader(\n",
        "                 dataset=train_set,\n",
        "                 batch_size=batch_size,\n",
        "                 shuffle=True)\n",
        "global_test_loader = torch.utils.data.DataLoader(\n",
        "                dataset=test_set,\n",
        "                batch_size=batch_size,\n",
        "                shuffle=False)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZjLFqabndU7W"
      },
      "source": [
        "# Define MNIST model\n",
        "class MLP(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(MLP, self).__init__()\n",
        "        self.fc1 = nn.Linear(28*28, 128)\n",
        "        self.fc2 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x.view(-1, 28*28)\n",
        "        x = self.fc1(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.fc2(x)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lGwFGzj-f03R"
      },
      "source": [
        "# Client Creation\n",
        "\n",
        "Next, let's create a set of 10 clients that logically represent out set of participants that are performing federated learning. Each client has its own partition of the training dataset and its own local model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YDAJZOAJf743"
      },
      "source": [
        "n_clients = 10\n",
        "\n",
        "def create_client(client_id, local_dataset, batch_size=32):\n",
        "  model =  MLP()\n",
        "  loader = torch.utils.data.DataLoader(\n",
        "                 dataset=local_dataset,\n",
        "                 batch_size=batch_size,\n",
        "                 shuffle=True)\n",
        "  return {\"client_id\": client_id,\n",
        "          \"local_dataset\": loader,\n",
        "          \"local_model\" : model,\n",
        "          \"optimizer\": optim.SGD(model.parameters(), lr=0.01, momentum=0.9)}\n",
        "\n",
        "# Partition datapoints\n",
        "local_datasets = [[] for i in range(n_clients)]\n",
        "for i, datapoint in enumerate(train_set):\n",
        "  local_datasets[i%n_clients].append(datapoint)\n",
        "\n",
        "# Create clients\n",
        "clients = [create_client(i, local_datasets[i]) for i in range(n_clients)]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ITU_UujnhwMi"
      },
      "source": [
        "# Federated Learning Training\n",
        "Now let's perform the federated learning loop! Recall that the steps of federated learning are:\n",
        "\n",
        "* Model broadcasting\n",
        "* Local training\n",
        "* Model aggregation\n",
        "* Model update\n",
        "\n",
        "The next code block will simulate this process with the controlling logic acting as the central server."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VFEUklZug0g7"
      },
      "source": [
        "def client_load_model(client, model):\n",
        "  client[\"local_model\"].load_state_dict(model.state_dict())\n",
        "\n",
        "def client_local_training(client):\n",
        "  criterion = nn.CrossEntropyLoss()\n",
        "  optimizer = client[\"optimizer\"]\n",
        "  dataset = client[\"local_dataset\"]\n",
        "  model = client[\"local_model\"]\n",
        "  for batch_idx, (x, target) in enumerate(dataset):\n",
        "    x, target = Variable(x), Variable(target)\n",
        "    out = model(x)\n",
        "    loss = criterion(out, target)\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()                \n",
        "\n",
        "def server_aggregate_models(clients):\n",
        "  models = [x[\"local_model\"].state_dict() for x in clients]\n",
        "  averaged_model = copy.deepcopy(models[0])\n",
        "  for k,v in averaged_model.items():\n",
        "    averaged_model[k] = sum([m[k] for m in models]) / len(models)\n",
        "  return averaged_model\n",
        "\n",
        "def evaluate(model, dataset):\n",
        "  criterion = nn.CrossEntropyLoss()\n",
        "  total_loss, total, correct = 0, 0, 0\n",
        "  for batch_idx, (x, target) in enumerate(dataset):\n",
        "    x, target = Variable(x), Variable(target)\n",
        "    out = model(x)\n",
        "    _, pred_label = torch.max(out.data, 1)\n",
        "    loss = criterion(out, target)\n",
        "    total_loss += loss.item()    \n",
        "    total += x.data.size()[0]\n",
        "    correct += (pred_label == target.data).sum()\n",
        "  print(\"Loss: %f, acc: %f\" % (total_loss, correct/total))\n",
        "\n",
        "def federated_learning_loop():  \n",
        "\n",
        "  global_model = MLP()\n",
        "  client_participation_fraction = .2\n",
        "  rounds = 1000\n",
        "\n",
        "  for r in range(rounds):\n",
        "\n",
        "    # Broadcast global model to clients\n",
        "    for c in clients:\n",
        "      client_load_model(c, global_model)\n",
        "    \n",
        "    # Selected clients perform local training\n",
        "    for c in clients:\n",
        "      if random.random() <= client_participation_fraction:\n",
        "        client_local_training(c)\n",
        "\n",
        "    # Aggregate the models\n",
        "    aggregated_model = server_aggregate_models(clients)\n",
        "\n",
        "    # Update global model\n",
        "    global_model.load_state_dict(aggregated_model)\n",
        "\n",
        "    # Evaluation\n",
        "    evaluate(global_model, global_test_loader)\n",
        "\n",
        "federated_learning_loop()\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}