2024/07/FederatedLearning/Part01/L2/L2_Federated_Training_Process.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9219af17",
   "metadata": {},
   "source": [
    "# Lesson 2: Federated Training Process"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59f3a5c4",
   "metadata": {},
   "source": [
    "Welcome to Lesson 2!\n",
    "\n",
    "To access the `requirements.txt` and `utils2.py` file for this course, go to `File` and click `Open`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c0ac7af",
   "metadata": {},
   "source": [
    "#### 1. Load imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb6907af-47cf-49c4-8c71-72f47bf99375",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flwr.client import Client, ClientApp, NumPyClient\n",
    "from flwr.common import ndarrays_to_parameters, Context\n",
    "from flwr.server import ServerApp, ServerConfig\n",
    "from flwr.server import ServerAppComponents\n",
    "from flwr.server.strategy import FedAvg\n",
    "from flwr.simulation import run_simulation\n",
    "\n",
    "from utils2 import *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b68fab28",
   "metadata": {},
   "source": [
    "#### 2. Set up the training and testing datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9991b9c",
   "metadata": {},
   "source": [
    "* Set up the MNIST training datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53299cb8-6e9b-4dac-8fa0-c620ef4a504c",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainset = datasets.MNIST(\n",
    "    \"./MNIST_data/\", download=True, train=True, transform=transform\n",
    ")\n",
    "\n",
    "total_length = len(trainset)\n",
    "split_size = total_length // 3\n",
    "torch.manual_seed(42)\n",
    "part1, part2, part3 = random_split(trainset, [split_size] * 3)\n",
    "\n",
    "part1 = exclude_digits(part1, excluded_digits=[1, 3, 7])\n",
    "part2 = exclude_digits(part2, excluded_digits=[2, 5, 8])\n",
    "part3 = exclude_digits(part3, excluded_digits=[4, 6, 9])\n",
    "\n",
    "train_sets = [part1, part2, part3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69561914",
   "metadata": {},
   "source": [
    "> Note: The function ```exclude_digits``` has been provided for your use. You can find it in the `utils2.py` file. \n",
    "\n",
    "  To access this file, go to `File` and click `Open`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e7c7bc5",
   "metadata": {},
   "source": [
    "* Set up the MNIST testing datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01e1aedf-5d70-44c9-a979-17b683a7219d",
   "metadata": {},
   "outputs": [],
   "source": [
    "testset = datasets.MNIST(\n",
    "    \"./MNIST_data/\", download=True, train=False, transform=transform\n",
    ")\n",
    "print(\"Number of examples in `testset`:\", len(testset))\n",
    "\n",
    "testset_137 = include_digits(testset, [1, 3, 7])\n",
    "testset_258 = include_digits(testset, [2, 5, 8])\n",
    "testset_469 = include_digits(testset, [4, 6, 9])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bac2f752",
   "metadata": {},
   "source": [
    "> Note: The function ```include_digits``` has been provided for your use. You can find it in the `utils2.py` file. \n",
    "\n",
    "  To access this file, go to `File` and click `Open`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aecd362f",
   "metadata": {},
   "source": [
    "#### 3. Define training and testing in the pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6fc6492",
   "metadata": {},
   "source": [
    "* Set functions for the client-server exchange of the training information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba2c96d-503c-4833-be91-c07f538dae48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sets the parameters of the model\n",
    "def set_weights(net, parameters):\n",
    "    params_dict = zip(net.state_dict().keys(), parameters)\n",
    "    state_dict = OrderedDict(\n",
    "        {k: torch.tensor(v) for k, v in params_dict}\n",
    "    )\n",
    "    net.load_state_dict(state_dict, strict=True)\n",
    "\n",
    "# Retrieves the parameters from the model\n",
    "def get_weights(net):\n",
    "    ndarrays = [\n",
    "        val.cpu().numpy() for _, val in net.state_dict().items()\n",
    "    ]\n",
    "    return ndarrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b5b4bfb",
   "metadata": {},
   "source": [
    "* Connect the training in the pipeline using the Flower Client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bab9677-8c5b-424d-a05f-b609bfe657d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FlowerClient(NumPyClient):\n",
    "    def __init__(self, net, trainset, testset):\n",
    "        self.net = net\n",
    "        self.trainset = trainset\n",
    "        self.testset = testset\n",
    "\n",
    "    # Train the model\n",
    "    def fit(self, parameters, config):\n",
    "        set_weights(self.net, parameters)\n",
    "        train_model(self.net, self.trainset)\n",
    "        return get_weights(self.net), len(self.trainset), {}\n",
    "\n",
    "    # Test the model\n",
    "    def evaluate(self, parameters: NDArrays, config: Dict[str, Scalar]):\n",
    "        set_weights(self.net, parameters)\n",
    "        loss, accuracy = evaluate_model(self.net, self.testset)\n",
    "        return loss, len(self.testset), {\"accuracy\": accuracy}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2050c06",
   "metadata": {},
   "source": [
    "* Flower calls `client_fn` whenever it needs an instance of one particular client to call fit or evaluate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "759eb41c-a7d8-4511-b0b5-7b114bf53f19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Client function\n",
    "def client_fn(context: Context) -> Client:\n",
    "    net = SimpleModel()\n",
    "    partition_id = int(context.node_config[\"partition-id\"])\n",
    "    client_train = train_sets[int(partition_id)]\n",
    "    client_test = testset\n",
    "    return FlowerClient(net, client_train, client_test).to_client()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2ca0734",
   "metadata": {},
   "source": [
    "* Create an instance of the ClientApp."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "659f4912-fcdf-41e3-aee1-b706c4f886cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = ClientApp(client_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "646dcc65",
   "metadata": {},
   "source": [
    "* Define `evaluate` for testing.\n",
    "\n",
    "The `evaluate` method evaluates the performance of the neural network model using the provided parameters and the test dataset (`testset`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fd32ef8-8c3c-4b45-b034-d1e57bb3b71f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(server_round, parameters, config):\n",
    "    net = SimpleModel()\n",
    "    set_weights(net, parameters)\n",
    "\n",
    "    _, accuracy = evaluate_model(net, testset)\n",
    "    _, accuracy137 = evaluate_model(net, testset_137)\n",
    "    _, accuracy258 = evaluate_model(net, testset_258)\n",
    "    _, accuracy469 = evaluate_model(net, testset_469)\n",
    "\n",
    "    log(INFO, \"test accuracy on all digits: %.4f\", accuracy)\n",
    "    log(INFO, \"test accuracy on [1,3,7]: %.4f\", accuracy137)\n",
    "    log(INFO, \"test accuracy on [2,5,8]: %.4f\", accuracy258)\n",
    "    log(INFO, \"test accuracy on [4,6,9]: %.4f\", accuracy469)\n",
    "\n",
    "    if server_round == 3:\n",
    "        cm = compute_confusion_matrix(net, testset)\n",
    "        plot_confusion_matrix(cm, \"Final Global Model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d15735c",
   "metadata": {},
   "source": [
    "* Strategy: Federated Average 'FedAvg'.\n",
    "\n",
    "The federated averaging strategy (`strategy.FedAvg`) is created for federated learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "497046c7-a944-4520-bccf-889d41802a0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = SimpleModel()\n",
    "params = ndarrays_to_parameters(get_weights(net))\n",
    "\n",
    "def server_fn(context: Context):\n",
    "    strategy = FedAvg(\n",
    "        fraction_fit=1.0,\n",
    "        fraction_evaluate=0.0,\n",
    "        initial_parameters=params,\n",
    "        evaluate_fn=evaluate,\n",
    "    )\n",
    "    config=ServerConfig(num_rounds=3)\n",
    "    return ServerAppComponents(\n",
    "        strategy=strategy,\n",
    "        config=config,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d44c734",
   "metadata": {},
   "source": [
    "* Create an instance of ServerApp."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bda25a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "server = ServerApp(server_fn=server_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16f5ee1c",
   "metadata": {},
   "source": [
    "* Start training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41a954c3-e564-4760-96d9-a8a7a879ca3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initiate the simulation passing the server and client apps\n",
    "# Specify the number of super nodes that will be selected on every round\n",
    "run_simulation(\n",
    "    server_app=server,\n",
    "    client_app=client,\n",
    "    num_supernodes=3,\n",
    "    backend_config=backend_setup,\n",
    ")"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}