examples/getting_started/pt/nvflare_pt_getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7a5c3d67-a6ea-4f59-84d2-effc3ef016e1",
   "metadata": {},
   "source": [
    "# Getting Started with NVFlare (PyTorch)\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NVIDIA/NVFlare/blob/main/examples/getting_started/pt/nvflare_pt_getting_started.ipynb)\n",
    "\n",
    "NVFlare is an open-source framework that allows researchers and\n",
    "data scientists to seamlessly move their machine learning and deep\n",
    "learning workflows into a federated paradigm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf2b4a8-ed42-421d-8898-c0c93f9d8a09",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Basic Concepts\n",
    "At the heart of NVFlare lies the concept of collaboration through\n",
    "\"tasks.\" An FL controller assigns tasks (e.g., training on local data) to one or more FL clients, processes returned\n",
    "results (e.g., model weight updates), and may assign additional\n",
    "tasks based on these results and other factors (e.g., a pre-configured\n",
    "number of training rounds). The clients run executors which can listen for tasks and perform the necessary computations locally, such as model training. This task-based interaction repeats\n",
    "until the experiment’s objectives are met. \n",
    "\n",
    "<img src=\"../../../docs/resources/controller_executor_no_filter.png\" alt=\"NVIDIA FLARE Controller and Executor\" width=75% height=75% />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "907933a8-20fd-4aa7-a3bf-3f5b5829a544",
   "metadata": {},
   "source": [
    "## Setup environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5496ffe",
   "metadata": {},
   "source": [
    "Install nvflare and dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33bba668-72ac-4e69-aaed-8d4254f547c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install --ignore-installed blinker\n",
    "! pip install nvflare~=2.5.0rc torch torchvision tensorboard"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baa517fe",
   "metadata": {},
   "source": [
    "If running in Google Colab, download the source code for this example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c139be6",
   "metadata": {},
   "outputs": [],
   "source": [
    "! npx degit NVIDIA/NVFlare/examples/getting_started/pt/src src"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b68cb248-dc6a-48d1-880d-33c4324d9723",
   "metadata": {},
   "source": [
    "## Federated Averaging with NVFlare\n",
    "Given the flexible controller and executor concepts, it is easy to implement different computing & communication patterns with NVFlare, such as [FedAvg](https://proceedings.mlr.press/v54/mcmahan17a?ref=https://githubhelp.com) and [cyclic weight transfer](https://academic.oup.com/jamia/article/25/8/945/4956468). \n",
    "\n",
    "The controller's `run()` routine is responsible for assigning tasks and processing task results from the Executors. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2f84fb1-9dd3-4c72-a727-c4614260f02f",
   "metadata": {},
   "source": [
    "### Server Code\n",
    "First, we provide a simple implementation of the [FedAvg](https://proceedings.mlr.press/v54/mcmahan17a?ref=https://githubhelp.com) algorithm with NVFlare. \n",
    "The `run()` routine implements the main algorithmic logic. \n",
    "Subroutines, like `sample_clients()` and `scatter_and_gather_model()` utilize the communicator object, native to each Controller to get the list of available clients,\n",
    "distribute the current global model to the clients, and collect their results.\n",
    "\n",
    "The FedAvg controller implements these main steps:\n",
    "1. FL server initializes an initial model using `self.load_model()`.\n",
    "2. For each round (global iteration):\n",
    "    - FL server samples available clients using `self.sample_clients()`.\n",
    "    - FL server sends the global model to clients and waits for their updates using `self.send_model_and_wait()`.\n",
    "    - FL server aggregates all the `results` and produces a new global model using `self.update_model()`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d62a13d5-1130-44e6-8818-70e30de401e6",
   "metadata": {},
   "source": [
    "```python\n",
    "class FedAvg(BaseFedAvg):\n",
    "    def run(self) -> None:\n",
    "        self.info(\"Start FedAvg.\")\n",
    "\n",
    "        model = self.load_model()\n",
    "        model.start_round = self.start_round\n",
    "        model.total_rounds = self.num_rounds\n",
    "\n",
    "        for self.current_round in range(self.start_round, self.start_round + self.num_rounds):\n",
    "            self.info(f\"Round {self.current_round} started.\")\n",
    "            model.current_round = self.current_round\n",
    "\n",
    "            clients = self.sample_clients(self.num_clients)\n",
    "\n",
    "            results = self.send_model_and_wait(targets=clients, data=model)\n",
    "\n",
    "            aggregate_results = self.aggregate(results)\n",
    "\n",
    "            model = self.update_model(model, aggregate_results)\n",
    "\n",
    "            self.save_model(model)\n",
    "\n",
    "        self.info(\"Finished FedAvg.\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d24b6476-089a-4e9d-825b-07107bd5d84a",
   "metadata": {},
   "source": [
    "### Client Code \n",
    "We take a CIFAR-10 example directly from [PyTorch website](https://github.com/pytorch/tutorials/blob/main/beginner_source/blitz/cifar10_tutorial.py) with some minor modifications, such as removing comments, move the network to [src/net.py](src/net.py), and add a main method and GPU support. The original code can be found at [cifar10_original.py](../../hello-world/ml-to-fl/pt/code/cifar10_original.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c551053-5460-4d83-8578-796074170342",
   "metadata": {},
   "source": [
    "Now, we need to adapt this centralized training code to something that can run in a federated setting.\n",
    "\n",
    "On the client side, the training workflow is as follows:\n",
    "1. Receive the model from the FL server.\n",
    "2. Perform local training on the received global model\n",
    "and/or evaluate the received global model for model\n",
    "selection.\n",
    "3. Send the new model back to the FL server."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c02bfc2a-783c-494f-9427-c38f40a2e870",
   "metadata": {},
   "source": [
    "Using NVFlare's client API, we can easily adapt machine learning code that was written for centralized training and apply it in a federated scenario.\n",
    "For a general use case, there are three essential methods to achieve this using the Client API :\n",
    "- `init()`: Initializes NVFlare Client API environment.\n",
    "- `receive()`: Receives model from the FL server.\n",
    "- `send()`: Sends the model to the FL server."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9115ee07-d848-4a7c-99ad-64e20ab7093c",
   "metadata": {},
   "source": [
    "With these simple methods, the developers can use the Client API\n",
    "to change their centralized training code to an FL scenario with\n",
    "five lines of code changes as shown below.\n",
    "```python\n",
    "    import nvflare.client as flare\n",
    "    \n",
    "    flare.init() # 1. Initializes NVFlare Client API environment.\n",
    "    input_model = flare.receive() # 2. Receives model from the FL server.\n",
    "    params = input_model.params # 3. Obtain the required information from the received model.\n",
    "    \n",
    "    # original local training code\n",
    "    new_params = local_train(params)\n",
    "    \n",
    "    output_model = flare.FLModel(params=new_params) # 4. Put the results in a new `FLModel`\n",
    "    flare.send(output_model) # 5. Sends the model to the FL server.  \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67432f44-4144-4347-8d74-e7f57e065a14",
   "metadata": {},
   "source": [
    "The full client training script is saved in a separate file, e.g. [./src/cifar10_fl.py](./src/cifar10_fl.py) doing CNN training on the [CIFAR-10](https://www.cs.toronto.edu/~kriz/cifar.html) dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da34414-bac4-4352-8077-ab7ade998eec",
   "metadata": {},
   "source": [
    "## Run an NVFlare Job\n",
    "Now that we have defined the FedAvg controller to run our federated compute workflow on the FL server, and our client training script to receive the global models, run local training, and send the results back to the FL server, we can put everything together using NVFlare's Job API."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cedaf75-3a4a-4843-8017-7716b53149a2",
   "metadata": {},
   "source": [
    "#### 1. Define the initial model\n",
    "First, we define the global model used to initialize the model on the FL server. See [src/net.py](src/net.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93889e62-b725-427c-8839-2771ca81d24c",
   "metadata": {},
   "source": [
    "```python\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, 5)\n",
    "        self.pool = nn.MaxPool2d(2, 2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = torch.flatten(x, 1)  # flatten all dimensions except batch\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b70da5d-ba8b-4e65-b47f-44bb9bddae4d",
   "metadata": {},
   "source": [
    "#### 2. Define a FedJob\n",
    "The `FedJob` is used to define how controllers and executors are placed within a federated job using the `to(object, target)` routine.\n",
    "\n",
    "Here we use a PyTorch `BaseFedJob`, where we can define the job name and the initial global model.\n",
    "The `BaseFedJob` automatically configures components for model persistence, model selection, and TensorBoard streaming for convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13771bfb-901f-485a-9a23-84db1ccd5fe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.net import Net\n",
    "\n",
    "from nvflare.app_common.workflows.fedavg import FedAvg\n",
    "from nvflare.app_opt.pt.job_config.base_fed_job import BaseFedJob\n",
    "from nvflare.job_config.script_runner import ScriptRunner\n",
    "\n",
    "job = BaseFedJob(\n",
    "    name=\"cifar10_pt_fedavg\",\n",
    "    initial_model=Net(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9361d9f8-54f3-4363-b3ba-706a7ae3a8e9",
   "metadata": {},
   "source": [
    "#### 3. Define the Controller Workflow\n",
    "Define the controller workflow and send to server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6962e6cc-995e-4356-8156-3ceba2c7a249",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_clients = 2\n",
    "\n",
    "controller = FedAvg(\n",
    "    num_clients=n_clients,\n",
    "    num_rounds=2,\n",
    ")\n",
    "job.to(controller, \"server\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77f5bc7f-4fb4-46e9-8f02-5e7245d95070",
   "metadata": {},
   "source": [
    "That completes the components that need to be defined on the server."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "548966c2-90bf-47ad-91d2-5c6c22c3c4f0",
   "metadata": {},
   "source": [
    "#### 4. Add clients\n",
    "Next, we can use the `ScriptRunner` and send it to each of the clients to run our training script.\n",
    "\n",
    "Note that our script could have additional input arguments, such as batch size or data path, but we don't use them here for simplicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad5d36fe-9ae5-43c3-80bc-2cdc66bf7a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(n_clients):\n",
    "    runner = ScriptRunner(\n",
    "        script=\"src/cifar10_fl.py\", script_args=\"\"  # f\"--batch_size 32 --data_path /tmp/data/site-{i}\"\n",
    "    )\n",
    "    job.to(runner, f\"site-{i+1}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "113fd6af-85be-4f75-8a8e-4666771252b3",
   "metadata": {},
   "source": [
    "That's it!\n",
    "\n",
    "#### 5. Optionally export the job\n",
    "Now, we could export the job and submit it to a real NVFlare deployment using the [Admin client](https://nvflare.readthedocs.io/en/main/real_world_fl/operation.html) or [FLARE API](https://nvflare.readthedocs.io/en/main/real_world_fl/flare_api.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99a270bf-c906-425b-b999-2306cb76eb62",
   "metadata": {},
   "outputs": [],
   "source": [
    "job.export_job(\"/tmp/nvflare/jobs/job_config\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac3f0a8-06bb-4bea-89d3-4a5fc5b76c63",
   "metadata": {},
   "source": [
    "#### 6. Run FL Simulation\n",
    "Finally, we can run our FedJob in simulation using NVFlare's [simulator](https://nvflare.readthedocs.io/en/main/user_guide/nvflare_cli/fl_simulator.html) under the hood. We can also specify which GPU should be used to run this client, which is helpful for simulated environments. The results will be saved in the specified `workdir`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13068ab7-35cf-49e7-91ed-10993049ef0d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "job.simulator_run(\"/tmp/nvflare/jobs/workdir\", gpu=\"0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ccbe893",
   "metadata": {},
   "source": [
    "If using Google Colab and the output is not showing correctly, export the job and run it with the simulator command instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e46a45d",
   "metadata": {},
   "outputs": [],
   "source": [
    "! nvflare simulator -w /tmp/nvflare/jobs/workdir -n 2 -t 2 -gpu 0 /tmp/nvflare/jobs/job_config/cifar10_pt_fedavg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f474ddfa-0d2e-4d7f-b033-8ccfbbd57a75",
   "metadata": {},
   "source": [
    "### Visualize the Training Results\n",
    "By default, we enable TensorBoard metric [streaming](https://nvflare.readthedocs.io/en/main/examples/tensorboard_streaming.html) using NVFlare's `SummaryWriter` in [src/cifar10_fl.py](src/cifar10_fl.py). \n",
    "\n",
    "The TensorBoard metrics will be received at the server, and you can visualize the training progress by running \n",
    "```commandline\n",
    "tensorboard --logdir=/tmp/nvflare/jobs/workdir/server/simulate_job/tb_events\n",
    "```\n",
    "in a new terminal.\n",
    "\n",
    "<img src=\"figs/tb_loss.png\" alt=\"Streamed TensorBoard metric\" width=50% height=50% />"
   ]
  }
 ],
 "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}