{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get started with Flower\n",
    "\n",
    "Welcome to the Flower federated learning tutorial!\n",
    "\n",
    "In this notebook, we'll build a federated learning system using the Flower framework, Flower Datasets and PyTorch. In part 1, we use PyTorch for the model training pipeline and data loading. In part 2, we federate the PyTorch project using Flower.\n",
    "\n",
    "> [Star Flower on GitHub](https://github.com/adap/flower) ⭐️ and join the Flower community on Flower Discuss and the Flower Slack to connect, ask questions, and get help:\n",
    "> - [Join Flower Discuss](https://discuss.flower.ai/) We'd love to hear from you in the `Introduction` topic! If anything is unclear, post in `Flower Help - Beginners`.\n",
    "> - [Join Flower Slack](https://flower.ai/join-slack) We'd love to hear from you in the `#introductions` channel! If anything is unclear, head over to the `#questions` channel.\n",
    "\n",
    "Let's get started! 🌼"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 0: Preparation\n",
    "\n",
    "Before we begin with any actual code, let's make sure that we have everything we need."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install dependencies\n",
    "\n",
    "Next, we install the necessary packages for PyTorch (`torch` and `torchvision`), Flower Datasets (`flwr-datasets`) and Flower (`flwr`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q flwr[simulation] flwr-datasets[vision] torch torchvision matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have all dependencies installed, we can import everything we need for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "from typing import List, Tuple\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "from datasets.utils.logging import disable_progress_bar\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import flwr\n",
    "from flwr.client import Client, ClientApp, NumPyClient\n",
    "from flwr.common import Metrics, Context\n",
    "from flwr.server import ServerApp, ServerConfig, ServerAppComponents\n",
    "from flwr.server.strategy import FedAvg\n",
    "from flwr.simulation import run_simulation\n",
    "from flwr_datasets import FederatedDataset\n",
    "\n",
    "DEVICE = torch.device(\"cpu\")  # Try \"cuda\" to train on GPU\n",
    "print(f\"Training on {DEVICE}\")\n",
    "print(f\"Flower {flwr.__version__} / PyTorch {torch.__version__}\")\n",
    "disable_progress_bar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is possible to switch to a runtime that has GPU acceleration enabled (on Google Colab: `Runtime > Change runtime type > Hardware accelerator: GPU > Save`). Note, however, that Google Colab is not always able to offer GPU acceleration. If you see an error related to GPU availability in one of the following sections, consider switching back to CPU-based execution by setting `DEVICE = torch.device(\"cpu\")`. If the runtime has GPU acceleration enabled, you should see the output `Training on cuda`, otherwise it'll say `Training on cpu`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the data\n",
    "\n",
    "Federated learning can be applied to many different types of tasks across different domains. In this tutorial, we introduce federated learning by training a simple convolutional neural network (CNN) on the popular CIFAR-10 dataset. CIFAR-10 can be used to train image classifiers that distinguish between images from ten different classes: 'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', and 'truck'."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We simulate having multiple datasets from multiple organizations (also called the \"cross-silo\" setting in federated learning) by splitting the original CIFAR-10 dataset into multiple partitions. Each partition will represent the data from a single organization. We're doing this purely for experimentation purposes, in the real world there's no need for data splitting because each organization already has their own data (the data is naturally partitioned).\n",
    "\n",
    "Each organization will act as a client in the federated learning system. Having ten organizations participate in a federation means having ten clients connected to the federated learning server.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the Flower Datasets library (`flwr-datasets`) to partition CIFAR-10 into ten partitions using `FederatedDataset`. We will create a small training and test set for each of the ten organizations and wrap each of these into a PyTorch `DataLoader`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLIENTS = 10\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "\n",
    "def load_datasets(partition_id: int):\n",
    "    fds = FederatedDataset(dataset=\"cifar10\", partitioners={\"train\": NUM_CLIENTS})\n",
    "    partition = fds.load_partition(partition_id)\n",
    "    # Divide data on each node: 80% train, 20% test\n",
    "    partition_train_test = partition.train_test_split(test_size=0.2, seed=42)\n",
    "    pytorch_transforms = transforms.Compose(\n",
    "        [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]\n",
    "    )\n",
    "\n",
    "    def apply_transforms(batch):\n",
    "        # Instead of passing transforms to CIFAR10(..., transform=transform)\n",
    "        # we will use this function to dataset.with_transform(apply_transforms)\n",
    "        # The transforms object is exactly the same\n",
    "        batch[\"img\"] = [pytorch_transforms(img) for img in batch[\"img\"]]\n",
    "        return batch\n",
    "\n",
    "    # Create train/val for each partition and wrap it into DataLoader\n",
    "    partition_train_test = partition_train_test.with_transform(apply_transforms)\n",
    "    trainloader = DataLoader(\n",
    "        partition_train_test[\"train\"], batch_size=BATCH_SIZE, shuffle=True\n",
    "    )\n",
    "    valloader = DataLoader(partition_train_test[\"test\"], batch_size=BATCH_SIZE)\n",
    "    testset = fds.load_split(\"test\").with_transform(apply_transforms)\n",
    "    testloader = DataLoader(testset, batch_size=BATCH_SIZE)\n",
    "    return trainloader, valloader, testloader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have a function that can return a training set and validation set (`trainloader` and `valloader`) representing one dataset from one of ten different organizations. Each `trainloader`/`valloader` pair contains 4000 training examples and 1000 validation examples. There's also a single `testloader` (we did not split the test set). Again, this is only necessary for building research or educational systems, actual federated learning systems have their data naturally distributed across multiple partitions.\n",
    "\n",
    "Let's take a look at the first batch of images and labels in the first training set (i.e., `trainloader` from `partition_id=0`) before we move on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainloader, _, _ = load_datasets(partition_id=0)\n",
    "batch = next(iter(trainloader))\n",
    "images, labels = batch[\"img\"], batch[\"label\"]\n",
    "\n",
    "# Reshape and convert images to a NumPy array\n",
    "# matplotlib requires images with the shape (height, width, 3)\n",
    "images = images.permute(0, 2, 3, 1).numpy()\n",
    "\n",
    "# Denormalize\n",
    "images = images / 2 + 0.5\n",
    "\n",
    "# Create a figure and a grid of subplots\n",
    "fig, axs = plt.subplots(4, 8, figsize=(12, 6))\n",
    "\n",
    "# Loop over the images and plot them\n",
    "for i, ax in enumerate(axs.flat):\n",
    "    ax.imshow(images[i])\n",
    "    ax.set_title(trainloader.dataset.features[\"label\"].int2str([labels[i]])[0])\n",
    "    ax.axis(\"off\")\n",
    "\n",
    "# Show the plot\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output above shows a random batch of images from the `trainloader` from the first of ten partitions. It also prints the labels associated with each image (i.e., one of the ten possible labels we've seen above). If you run the cell again, you should see another batch of images."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Centralized Training with PyTorch\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we're going to use PyTorch to define a simple convolutional neural network. This introduction assumes basic familiarity with PyTorch, so it doesn't cover the PyTorch-related aspects in full detail. If you want to dive deeper into PyTorch, we recommend [*DEEP LEARNING WITH PYTORCH: A 60 MINUTE BLITZ*](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the model\n",
    "\n",
    "We use the simple CNN described in the [PyTorch tutorial](https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#define-a-convolutional-neural-network):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self) -> None:\n",
    "        super(Net, self).__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: torch.Tensor) -> torch.Tensor:\n",
    "        x = self.pool(F.relu(self.conv1(x)))\n",
    "        x = self.pool(F.relu(self.conv2(x)))\n",
    "        x = x.view(-1, 16 * 5 * 5)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's continue with the usual training and test functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(net, trainloader, epochs: int, verbose=False):\n",
    "    \"\"\"Train the network on the training set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(net.parameters())\n",
    "    net.train()\n",
    "    for epoch in range(epochs):\n",
    "        correct, total, epoch_loss = 0, 0, 0.0\n",
    "        for batch in trainloader:\n",
    "            images, labels = batch[\"img\"].to(DEVICE), batch[\"label\"].to(DEVICE)\n",
    "            optimizer.zero_grad()\n",
    "            outputs = net(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            # Metrics\n",
    "            epoch_loss += loss\n",
    "            total += labels.size(0)\n",
    "            correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()\n",
    "        epoch_loss /= len(trainloader.dataset)\n",
    "        epoch_acc = correct / total\n",
    "        if verbose:\n",
    "            print(f\"Epoch {epoch+1}: train loss {epoch_loss}, accuracy {epoch_acc}\")\n",
    "\n",
    "\n",
    "def test(net, testloader):\n",
    "    \"\"\"Evaluate the network on the entire test set.\"\"\"\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    correct, total, loss = 0, 0, 0.0\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in testloader:\n",
    "            images, labels = batch[\"img\"].to(DEVICE), batch[\"label\"].to(DEVICE)\n",
    "            outputs = net(images)\n",
    "            loss += criterion(outputs, labels).item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    loss /= len(testloader.dataset)\n",
    "    accuracy = correct / total\n",
    "    return loss, accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the model\n",
    "\n",
    "We now have all the basic building blocks we need: a dataset, a model, a training function, and a test function. Let's put them together to train the model on the dataset of one of our organizations (`partition_id=0`). This simulates the reality of most machine learning projects today: each organization has their own data and trains models only on this internal data: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainloader, valloader, testloader = load_datasets(partition_id=0)\n",
    "net = Net().to(DEVICE)\n",
    "\n",
    "for epoch in range(5):\n",
    "    train(net, trainloader, 1)\n",
    "    loss, accuracy = test(net, valloader)\n",
    "    print(f\"Epoch {epoch+1}: validation loss {loss}, accuracy {accuracy}\")\n",
    "\n",
    "loss, accuracy = test(net, testloader)\n",
    "print(f\"Final test set performance:\\n\\tloss {loss}\\n\\taccuracy {accuracy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training the simple CNN on our CIFAR-10 split for 5 epochs should result in a test set accuracy of about 41%, which is not good, but at the same time, it doesn't really matter for the purposes of this tutorial. The intent was just to show a simple centralized training pipeline that sets the stage for what comes next - federated learning!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Federated Learning with Flower\n",
    "\n",
    "Step 1 demonstrated a simple centralized training pipeline. All data was in one place (i.e., a single `trainloader` and a single `valloader`). Next, we'll simulate a situation where we have multiple datasets in multiple organizations and where we train a model over these organizations using federated learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Update model parameters\n",
    "\n",
    "In federated learning, the server sends global model parameters to the client, and the client updates the local model with parameters received from the server. It then trains the model on the local data (which changes the model parameters locally) and sends the updated/changed model parameters back to the server (or, alternatively, it sends just the gradients back to the server, not the full model parameters).\n",
    "\n",
    "We need two helper functions to update the local model with parameters received from the server and to get the updated model parameters from the local model: `set_parameters` and `get_parameters`. The following two functions do just that for the PyTorch model above.\n",
    "\n",
    "The details of how this works are not really important here (feel free to consult the PyTorch documentation if you want to learn more). In essence, we use `state_dict` to access PyTorch model parameter tensors. The parameter tensors are then converted to/from a list of NumPy ndarray's (which the Flower `NumPyClient` knows how to serialize/deserialize):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_parameters(net, parameters: List[np.ndarray]):\n",
    "    params_dict = zip(net.state_dict().keys(), parameters)\n",
    "    state_dict = OrderedDict({k: torch.Tensor(v) for k, v in params_dict})\n",
    "    net.load_state_dict(state_dict, strict=True)\n",
    "\n",
    "\n",
    "def get_parameters(net) -> List[np.ndarray]:\n",
    "    return [val.cpu().numpy() for _, val in net.state_dict().items()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Flower ClientApp\n",
    "\n",
    "With that out of the way, let's move on to the interesting part. Federated learning systems consist of a server and multiple clients. In Flower, we create a `ServerApp` and a `ClientApp` to run the server-side and client-side code, respectively.\n",
    "\n",
    "The first step toward creating a `ClientApp` is to implement a subclasses of `flwr.client.Client` or `flwr.client.NumPyClient`. We use `NumPyClient` in this tutorial because it is easier to implement and requires us to write less boilerplate. To implement `NumPyClient`, we create a subclass that implements the three methods `get_parameters`, `fit`, and `evaluate`:\n",
    "\n",
    "* `get_parameters`: Return the current local model parameters\n",
    "* `fit`: Receive model parameters from the server, train the model on the local data, and return the updated model parameters to the server\n",
    "* `evaluate`: Receive model parameters from the server, evaluate the model on the local data, and return the evaluation result to the server\n",
    "\n",
    "We mentioned that our clients will use the previously defined PyTorch components for model training and evaluation. Let's see a simple Flower client implementation that brings everything together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FlowerClient(NumPyClient):\n",
    "    def __init__(self, net, trainloader, valloader):\n",
    "        self.net = net\n",
    "        self.trainloader = trainloader\n",
    "        self.valloader = valloader\n",
    "\n",
    "    def get_parameters(self, config):\n",
    "        return get_parameters(self.net)\n",
    "\n",
    "    def fit(self, parameters, config):\n",
    "        set_parameters(self.net, parameters)\n",
    "        train(self.net, self.trainloader, epochs=1)\n",
    "        return get_parameters(self.net), len(self.trainloader), {}\n",
    "\n",
    "    def evaluate(self, parameters, config):\n",
    "        set_parameters(self.net, parameters)\n",
    "        loss, accuracy = test(self.net, self.valloader)\n",
    "        return float(loss), len(self.valloader), {\"accuracy\": float(accuracy)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our class `FlowerClient` defines how local training/evaluation will be performed and allows Flower to call the local training/evaluation through `fit` and `evaluate`. Each instance of `FlowerClient` represents a *single client* in our federated learning system. Federated learning systems have multiple clients (otherwise, there's not much to federate), so each client will be represented by its own instance of `FlowerClient`. If we have, for example, three clients in our workload, then we'd have three instances of `FlowerClient` (one on each of the machines we'd start the client on). Flower calls `FlowerClient.fit` on the respective instance when the server selects a particular client for training (and `FlowerClient.evaluate` for evaluation).\n",
    "\n",
    "In this notebook, we want to simulate a federated learning system with 10 clients *on a single machine*. This means that the server and all 10 clients will live on a single machine and share resources such as CPU, GPU, and memory. Having 10 clients would mean having 10 instances of `FlowerClient` in memory. Doing this on a single machine can quickly exhaust the available memory resources, even if only a subset of these clients participates in a single round of federated learning.\n",
    "\n",
    "In addition to the regular capabilities where server and clients run on multiple machines, Flower, therefore, provides special simulation capabilities that create `FlowerClient` instances only when they are actually necessary for training or evaluation. To enable the Flower framework to create clients when necessary, we need to implement a function that creates a `FlowerClient` instance on demand. We typically call this function `client_fn`. Flower calls `client_fn` whenever it needs an instance of one particular client to call `fit` or `evaluate` (those instances are usually discarded after use, so they should not keep any local state). In federated learning experiments using Flower, clients are identified by a partition ID, or `partition-id`. This `partition-id` is used to load different local data partitions for different clients, as can be seen below. The value of `partition-id` is retrieved from the `node_config` dictionary in the `Context` object, which holds the information that persists throughout each training round. \n",
    "\n",
    "With this, we have the class `FlowerClient` which defines client-side training/evaluation and `client_fn` which allows Flower to create `FlowerClient` instances whenever it needs to call `fit` or `evaluate` on one particular client. Last, but definitely not least, we create an instance of `ClientApp` and pass it the `client_fn`. `ClientApp` is the entrypoint that a running Flower client uses to call your code (as defined in, for example, `FlowerClient.fit`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def client_fn(context: Context) -> Client:\n",
    "    \"\"\"Create a Flower client representing a single organization.\"\"\"\n",
    "\n",
    "    # Load model\n",
    "    net = Net().to(DEVICE)\n",
    "\n",
    "    # Load data (CIFAR-10)\n",
    "    # Note: each client gets a different trainloader/valloader, so each client\n",
    "    # will train and evaluate on their own unique data partition\n",
    "    # Read the node_config to fetch data partition associated to this node\n",
    "    partition_id = context.node_config[\"partition-id\"]\n",
    "    trainloader, valloader, _ = load_datasets(partition_id=partition_id)\n",
    "\n",
    "    # Create a single Flower client representing a single organization\n",
    "    # FlowerClient is a subclass of NumPyClient, so we need to call .to_client()\n",
    "    # to convert it to a subclass of `flwr.client.Client`\n",
    "    return FlowerClient(net, trainloader, valloader).to_client()\n",
    "\n",
    "\n",
    "# Create the ClientApp\n",
    "client = ClientApp(client_fn=client_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Flower ServerApp\n",
    "\n",
    "On the server side, we need to configure a strategy which encapsulates the federated learning approach/algorithm, for example, *Federated Averaging* (FedAvg). Flower has a number of built-in strategies, but we can also use our own strategy implementations to customize nearly all aspects of the federated learning approach. For this example, we use the built-in `FedAvg` implementation and customize it using a few basic parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create FedAvg strategy\n",
    "strategy = FedAvg(\n",
    "    fraction_fit=1.0,  # Sample 100% of available clients for training\n",
    "    fraction_evaluate=0.5,  # Sample 50% of available clients for evaluation\n",
    "    min_fit_clients=10,  # Never sample less than 10 clients for training\n",
    "    min_evaluate_clients=5,  # Never sample less than 5 clients for evaluation\n",
    "    min_available_clients=10,  # Wait until all 10 clients are available\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similar to `ClientApp`, we create a `ServerApp` using a utility function `server_fn`. In `server_fn`, we pass an instance of `ServerConfig` for defining the number of federated learning rounds (`num_rounds`) and we also pass the previously created `strategy`. The `server_fn` returns a `ServerAppComponents` object containing the settings that define the `ServerApp` behaviour. `ServerApp` is the entrypoint that Flower uses to call all your server-side code (for example, the strategy)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def server_fn(context: Context) -> ServerAppComponents:\n",
    "    \"\"\"Construct components that set the ServerApp behaviour.\n",
    "\n",
    "    You can use the settings in `context.run_config` to parameterize the\n",
    "    construction of all elements (e.g the strategy or the number of rounds)\n",
    "    wrapped in the returned ServerAppComponents object.\n",
    "    \"\"\"\n",
    "\n",
    "    # Configure the server for 5 rounds of training\n",
    "    config = ServerConfig(num_rounds=5)\n",
    "\n",
    "    return ServerAppComponents(strategy=strategy, config=config)\n",
    "\n",
    "\n",
    "# Create the ServerApp\n",
    "server = ServerApp(server_fn=server_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the training\n",
    "\n",
    "In simulation, we often want to control the amount of resources each client can use. In the next cell, we specify a `backend_config` dictionary with the `client_resources` key (required) for defining the amount of CPU and GPU resources each client can access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the resources each of your clients need\n",
    "# By default, each client will be allocated 1x CPU and 0x GPUs\n",
    "backend_config = {\"client_resources\": {\"num_cpus\": 1, \"num_gpus\": 0.0}}\n",
    "\n",
    "# When running on GPU, assign an entire GPU for each client\n",
    "if DEVICE.type == \"cuda\":\n",
    "    backend_config = {\"client_resources\": {\"num_cpus\": 1, \"num_gpus\": 1.0}}\n",
    "    # Refer to our Flower framework documentation for more details about Flower simulations\n",
    "    # and how to set up the `backend_config`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last step is the actual call to `run_simulation` which - you guessed it - runs the simulation. `run_simulation` accepts a number of arguments:\n",
    "- `server_app` and `client_app`: the previously created `ServerApp` and `ClientApp` objects, respectively\n",
    "- `num_supernodes`: the number of `SuperNodes` to simulate which equals the number of clients for Flower simulation\n",
    "- `backend_config`: the resource allocation used in this simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run simulation\n",
    "run_simulation(\n",
    "    server_app=server,\n",
    "    client_app=client,\n",
    "    num_supernodes=NUM_CLIENTS,\n",
    "    backend_config=backend_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Behind the scenes\n",
    "\n",
    "So how does this work? How does Flower execute this simulation?\n",
    "\n",
    "When we call `run_simulation`, we tell Flower that there are 10 clients (`num_supernodes=10`, where 1 `SuperNode` launches 1 `ClientApp`). Flower then goes ahead an asks the `ServerApp` to issue an instructions to those nodes using the `FedAvg` strategy. `FedAvg` knows that it should select 100% of the available clients (`fraction_fit=1.0`), so it goes ahead and selects 10 random clients (i.e., 100% of 10).\n",
    "\n",
    "Flower then asks the selected 10 clients to train the model. Each of the 10 `ClientApp` instances receives a message, which causes it to call `client_fn` to create an instance of `FlowerClient`. It then calls `.fit()` on each the `FlowerClient` instances and returns the resulting model parameter updates to the `ServerApp`. When the `ServerApp` receives the model parameter updates from the clients, it hands those updates over to the strategy (*FedAvg*) for aggregation. The strategy aggregates those updates and returns the new global model, which then gets used in the next round of federated learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Where's the accuracy?\n",
    "\n",
    "You may have noticed that all metrics except for `losses_distributed` are empty. Where did the `{\"accuracy\": float(accuracy)}` go?\n",
    "\n",
    "Flower can automatically aggregate losses returned by individual clients, but it cannot do the same for metrics in the generic metrics dictionary (the one with the `accuracy` key). Metrics dictionaries can contain very different kinds of metrics and even key/value pairs that are not metrics at all, so the framework does not (and can not) know how to handle these automatically.\n",
    "\n",
    "As users, we need to tell the framework how to handle/aggregate these custom metrics, and we do so by passing metric aggregation functions to the strategy. The strategy will then call these functions whenever it receives fit or evaluate metrics from clients. The two possible functions are `fit_metrics_aggregation_fn` and `evaluate_metrics_aggregation_fn`.\n",
    "\n",
    "Let's create a simple weighted averaging function to aggregate the `accuracy` metric we return from `evaluate`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:\n",
    "    # Multiply accuracy of each client by number of examples used\n",
    "    accuracies = [num_examples * m[\"accuracy\"] for num_examples, m in metrics]\n",
    "    examples = [num_examples for num_examples, _ in metrics]\n",
    "\n",
    "    # Aggregate and return custom metric (weighted average)\n",
    "    return {\"accuracy\": sum(accuracies) / sum(examples)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def server_fn(context: Context) -> ServerAppComponents:\n",
    "    \"\"\"Construct components that set the ServerApp behaviour.\n",
    "\n",
    "    You can use settings in `context.run_config` to parameterize the\n",
    "    construction of all elements (e.g the strategy or the number of rounds)\n",
    "    wrapped in the returned ServerAppComponents object.\n",
    "    \"\"\"\n",
    "\n",
    "    # Create FedAvg strategy\n",
    "    strategy = FedAvg(\n",
    "        fraction_fit=1.0,\n",
    "        fraction_evaluate=0.5,\n",
    "        min_fit_clients=10,\n",
    "        min_evaluate_clients=5,\n",
    "        min_available_clients=10,\n",
    "        evaluate_metrics_aggregation_fn=weighted_average,  # <-- pass the metric aggregation function\n",
    "    )\n",
    "\n",
    "    # Configure the server for 5 rounds of training\n",
    "    config = ServerConfig(num_rounds=5)\n",
    "\n",
    "    return ServerAppComponents(strategy=strategy, config=config)\n",
    "\n",
    "\n",
    "# Create a new server instance with the updated FedAvg strategy\n",
    "server = ServerApp(server_fn=server_fn)\n",
    "\n",
    "# Run simulation\n",
    "run_simulation(\n",
    "    server_app=server,\n",
    "    client_app=client,\n",
    "    num_supernodes=NUM_CLIENTS,\n",
    "    backend_config=backend_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have a full system that performs federated training and federated evaluation. It uses the `weighted_average` function to aggregate custom evaluation metrics and calculates a single `accuracy` metric across all clients on the server side.\n",
    "\n",
    "The other two categories of metrics (`losses_centralized` and `metrics_centralized`) are still empty because they only apply when centralized evaluation is being used. Part two of the Flower tutorial will cover centralized evaluation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final remarks\n",
    "\n",
    "Congratulations, you just trained a convolutional neural network, federated over 10 clients! With that, you understand the basics of federated learning with Flower. The same approach you've seen can be used with other machine learning frameworks (not just PyTorch) and tasks (not just CIFAR-10 images classification), for example NLP with Hugging Face Transformers or speech with SpeechBrain.\n",
    "\n",
    "In the next notebook, we're going to cover some more advanced concepts. Want to customize your strategy? Initialize parameters on the server side? Or evaluate the aggregated model on the server side? We'll cover all this and more in the next tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "Before you continue, make sure to join the Flower community on Flower Discuss ([Join Flower Discuss](https://discuss.flower.ai)) and on Slack ([Join Slack](https://flower.ai/join-slack/)).\n",
    "\n",
    "There's a dedicated `#questions` channel if you need help, but we'd also love to hear who you are in `#introductions`!\n",
    "\n",
    "The [Flower Federated Learning Tutorial - Part 2](https://flower.ai/docs/framework/tutorial-use-a-federated-learning-strategy-pytorch.html) goes into more depth about strategies and all the advanced things you can build with them.\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Flower-1-Intro-to-FL-PyTorch.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}