Implement FedAvg from Scratch in PyTorch

Build Federated Averaging from scratch in PyTorch with IID vs non-IID comparison.

Introduction

30 Second Summary

Hospitals, banks, and research labs hold sensitive data they cannot share with each other. They still want to train a shared AI model without ever moving the data off their own servers.

In this project, you will implement the Federated Averaging (FedAvg) algorithm from scratch in PyTorch. You will simulate 10 clients training a model on their own local data, then averaging the results into a single improved global model.

What You'll Build

A from-scratch FedAvg implementation that trains a CNN across 10 simulated clients on the MNIST dataset, producing convergence comparison plots and a test suite.

By the end of this project, you'll have:

  • A CNN model matching the McMahan et al. 2017 paper architecture, with a Dirichlet distribution partitioner simulating realistic non-IID federated data across 10 clients.
  • A complete FedAvg training loop with weighted parameter aggregation across 10 simulated clients and side-by-side convergence plots comparing IID vs non-IID performance.
  • A pytest test suite with a convergence assertion proving your implementation works, plus a research-audience README with results table and architecture diagram.
  • Secret Mission: Implement the Krum Byzantine-resilient aggregation algorithm, simulate a malicious client, and prove Krum resists the attack while FedAvg collapses.

Are there any prerequisites?

You need basic Python (functions, loops, dictionaries) and some experience with TensorFlow or Keras.

No prior Federated Learning or PyTorch knowledge is required. This project teaches both from the ground up.

Before We Start

Before you write any Federated Learning code, you need a clean project environment with the right tools installed. This step gets your machine ready so you can focus entirely on building the FedAvg algorithm.

You will create an isolated Python virtual environment, install PyTorch and supporting libraries, initialize a Git repository, and lay out the file structure that mirrors how real FL systems are architected.

In this step, get ready to:
  • Set up the project directory and Python virtual environment.
  • Install PyTorch, matplotlib, and pytest.
  • Create the file structure for the FedAvg implementation.
Create your project directory and virtual environment
  • Press the Windows key to open your search bar.
  • Type Antigravity and press Enter to open Google Antigravity IDE.
  • Open the integrated terminal by pressing Ctrl+`.

A virtual environment is an isolated Python installation that keeps this project's packages separate from your system Python. You will create one inside your project folder.

  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder, move into it, and set up the virtual environment by running these commands:
mkdir fedavg-from-scratch
cd fedavg-from-scratch
python -m venv venv
.\venv\Scripts\Activate.ps1

What does this do?

  • mkdir fedavg-from-scratch creates a new folder for the project.
  • python -m venv venv creates a virtual environment named venv inside the project folder.
  • .\venv\Scripts\Activate.ps1 activates the virtual environment so that any packages you install go into this project only.
  • Confirm the virtual environment is active by looking at your terminal prompt.

You should see (venv) at the beginning of your terminal prompt. This confirms all packages you install will stay inside this project.

Don't see (venv) in your prompt?

  • If you get a "running scripts is disabled" error, PowerShell's execution policy is blocking the activation script. Run Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser and then try the activate command again.
  • Make sure you are inside the fedavg-from-scratch folder before running the activate command.

Help me activate my Python virtual environment in Antigravity IDE on Windows.

  • Check your Python version by running this command:
python --version

✔️ I see Python 3.11 or higher

Your Python version is compatible. You are ready to install the project dependencies.

ⓧ I see Python 3.10

The project still works on Python 3.10. When you install packages in the next substep, pip will automatically select compatible versions of matplotlib and numpy for your Python version. You can proceed with the next substep.

ⓧ I see Python 3.9 or lower, or an error

PyTorch requires Python 3.10 or higher. You need to install a newer version of Python before continuing.

  • Download and install Python 3.11 or higher from python.org/downloads.
  • During installation, check the box that says Add Python to PATH.
  • After installation, close and reopen the terminal in Antigravity IDE, then re-run the commands from the beginning of this substep.
  • Open the project folder in Antigravity IDE by clicking File in the top menu bar.
  • Click Open Folder.
  • Select the fedavg-from-scratch folder on your Desktop, then click Select Folder.

Why open the folder?

Opening the folder gives Antigravity IDE access to your project's file tree in the Explorer sidebar. You will use this sidebar to create and edit files throughout the project.

If Antigravity IDE opens a new window, reopen the terminal with Ctrl+` and re-activate your virtual environment by running:

.\venv\Scripts\Activate.ps1
Install dependencies and initialize Git

Your project depends on three groups of packages: PyTorch for building and training neural networks, matplotlib for plotting convergence curves, and pytest for verifying your implementation with automated tests.

  • Create a new file called requirements.txt by right-clicking in the Explorer sidebar and selecting New File.
  • Name the file requirements.txt and paste the following content:
torch
torchvision
torchaudio
matplotlib
pytest

What is requirements.txt?

This file documents which packages the project depends on. It is a standard Python convention. Anyone who clones your repo can install everything with one command.

  • Save requirements.txt by pressing Ctrl+S.

You will install PyTorch separately from the other packages because PyTorch needs a special download URL to get the CPU-only build.

  • Install PyTorch, torchvision, and torchaudio by running this command in your terminal:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

What does --index-url do?

The --index-url flag tells pip to download PyTorch from a specific repository that contains the CPU-only build. This version is smaller and sufficient for MNIST training.

Have an NVIDIA GPU?

If you have an NVIDIA GPU and want faster training, replace cpu with cu128 (for CUDA 12.8) or cu118 (for CUDA 11.8) in the URL. The CPU build works fine for this project.

  • Install matplotlib and pytest by running:
pip install matplotlib pytest
  • Verify PyTorch installed correctly by running:
python -c "import torch; print(torch.__version__)"

You should see a version string printed to the terminal (e.g. 2.12.1+cpu). The exact number may differ. As long as a version prints without an error, PyTorch is ready.

Seeing a ModuleNotFoundError?

  • Make sure (venv) is still visible in your terminal prompt. If it disappeared, re-activate with .\venv\Scripts\Activate.ps1 and run the install commands again.
  • If the install itself failed, check your internet connection and try running the pip install command again.

Help me fix my PyTorch installation in my virtual environment.

Why PyTorch instead of TensorFlow?

Your entire BOPFL research stack is PyTorch-based. Opacus (for differential privacy), HuggingFace PEFT (for LoRA), and every paper on your reading list use PyTorch.

Learning PyTorch now means you will not need to context-switch when you reach Projects 2 through 6.

  • Initialize a Git repository by running:
git init

Your project folder is now a Git repository. Next, you will tell Git which files to ignore.

  • Create a new file called .gitignore by right-clicking in the Explorer sidebar and selecting New File.
  • Name the file .gitignore and paste the following content:
venv/
data/
results/
__pycache__/
*.pyc
.pytest_cache/

What does .gitignore do?

Any file or folder you list here will be completely invisible to Git. This keeps generated data, the virtual environment, and Python cache files out of your repository.

  • Save .gitignore by pressing Ctrl+S.

✔️ Awesome, I've got everything!

Great. Make sure both requirements.txt and .gitignore are saved before moving on.

ⓧ I'd like to double check the full code

torch
torchvision
torchaudio
matplotlib
pytest
venv/
data/
results/
__pycache__/
*.pyc
.pytest_cache/
Create the file structure

Your project needs five Python files and three directories. Each file has a single responsibility, mirroring how real FL systems separate clients from servers.

  • Create all the project files and directories by running these commands in your terminal:
New-Item model.py
New-Item data.py
New-Item client.py
New-Item server.py
New-Item train.py
mkdir tests
New-Item tests\test_fedavg.py
mkdir results
mkdir data

What is each file for?

  • model.py defines the neural network architecture.
  • data.py splits data across simulated clients using a Dirichlet distribution.
  • client.py handles local training on each client.
  • server.py handles aggregation of client updates into a global model.
  • train.py orchestrates the full experiment and produces comparison plots.
  • tests/test_fedavg.py contains automated tests to verify your implementation is correct.
  • Confirm the files exist by running:
ls *.py

You should see client.py, data.py, model.py, server.py, and train.py listed in the output.

Not seeing all five files?

  • Make sure your terminal is in the fedavg-from-scratch folder. Run pwd to check your current location.
  • If a file is missing, create it individually with New-Item filename.py, replacing filename.py with the missing file name.

Help me create the project file structure for my FedAvg project.

Your environment is ready. Next up, you will build the CNN model and the Dirichlet data partitioner that simulates how hospitals would each hold different slices of patient data.

Build the CNN Model and Dirichlet Data Partitioner

Your environment is set up and all dependencies are installed. This step introduces two foundational pieces of any federated learning system: the neural network that each client trains, and the data partitioner that simulates how data is distributed across clients.

If you have only used TensorFlow or Keras before, this is where you learn the PyTorch way of doing things. You will also learn why data heterogeneity is the central challenge in federated learning.

In this step, get ready to:
  • Build the CNN architecture from the McMahan et al. 2017 paper in PyTorch.
  • Implement a Dirichlet distribution data partitioner that simulates realistic federated data splits.
Implement the MNIST CNN

In PyTorch, you define a model by creating a class that inherits from nn.Module. You declare the layers in __init__ and define how data flows through them in forward.

The CNN you are building matches the exact architecture from the original FedAvg paper. It has 2 convolutional layers (32 and 64 filters, 5x5 kernels) each followed by 2x2 max pooling, then 2 fully connected layers. The total is 1,663,370 parameters.

  • Open model.py in the Explorer sidebar and paste the following code:
import torch.nn as nn


class MNISTCnn(nn.Module):
    """CNN for MNIST classification following McMahan et al. 2017.

    Architecture:
        Conv(1->32, 5x5, pad=2) -> ReLU -> MaxPool(2x2)
        Conv(32->64, 5x5, pad=2) -> ReLU -> MaxPool(2x2)
        Linear(64*7*7 -> 512) -> ReLU
        Linear(512 -> 10)

    Total parameters: 1,663,370
    """

    def __init__(self):
        super().__init__()
        # Two convolutional layers with increasing filter count
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5, padding=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, padding=2)
        # 2x2 max pooling halves spatial dimensions each time
        self.pool = nn.MaxPool2d(2, 2)
        # Fully connected layers: flatten -> 512 -> 10 output classes
        self.fc1 = nn.Linear(64 * 7 * 7, 512)
        self.fc2 = nn.Linear(512, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        # Conv1 + ReLU + Pool: 28x28 -> 14x14
        x = self.pool(self.relu(self.conv1(x)))
        # Conv2 + ReLU + Pool: 14x14 -> 7x7
        x = self.pool(self.relu(self.conv2(x)))
        # Flatten 64 feature maps of 7x7 into a vector of 3136
        x = x.view(-1, 64 * 7 * 7)
        # FC layers with ReLU, then raw logits for 10 digit classes
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

How do the dimensions work?

MNIST images are 28x28 pixels with 1 grayscale channel. After the first convolution and pooling, they become 14x14. After the second, 7x7.

You then flatten 64 feature maps of 7x7 into a single vector of 3,136 values (64 * 7 * 7). This feeds into the fully connected layers, which narrow down to 10 outputs (one per digit class). Understanding this spatial transformation is fundamental PyTorch knowledge.

  • Save model.py.
  • Verify the model has the correct parameter count by running this command in your terminal:
python -c "from model import MNISTCnn; m = MNISTCnn(); print(sum(p.numel() for p in m.parameters()))"
  • You should see 1663370 printed in the terminal. This confirms your architecture matches the paper exactly.

Seeing an ImportError?

  • Make sure you saved model.py in the project root (fedavg-from-scratch), not in a subfolder.
  • Confirm your terminal is in the fedavg-from-scratch directory and the (venv) prefix is visible.
  • Check that the class name is exactly MNISTCnn (capital M, capital C).

Why can't Python find my MNISTCnn class?

Why this specific architecture?

This is the exact CNN from the original FedAvg paper (McMahan et al. 2017). Using the same model lets you compare your results against published numbers. The architecture is also small enough to train on CPU in minutes, which keeps this project under an hour.

✔️ Awesome, I've got everything!

Great. Double check you have saved model.py and confirmed the parameter count is 1663370.

ⓧ I'd like to double check the full code

import torch.nn as nn


class MNISTCnn(nn.Module):
    """CNN for MNIST classification following McMahan et al. 2017.

    Architecture:
        Conv(1->32, 5x5, pad=2) -> ReLU -> MaxPool(2x2)
        Conv(32->64, 5x5, pad=2) -> ReLU -> MaxPool(2x2)
        Linear(64*7*7 -> 512) -> ReLU
        Linear(512 -> 10)

    Total parameters: 1,663,370
    """

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5, padding=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, padding=2)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 512)
        self.fc2 = nn.Linear(512, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 7 * 7)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x
Implement Dirichlet data partitioning

In real federated learning, each client has its own data that looks different from other clients. One hospital might see mostly heart disease patients while another sees mostly broken bones. This is called non-IID data (non-Independent and Identically Distributed).

To simulate this, you use a Dirichlet distribution. Think of it as a "fairness dial" controlled by a single parameter called alpha. High alpha (like 100) splits data almost evenly across clients. Low alpha (like 0.1) creates dramatic imbalance where some clients hold mostly a single digit class.

What is IID vs non-IID?

IID means every client sees a balanced, representative sample of all the data. If you shuffled all MNIST digits into a deck and dealt them out evenly, each client would get roughly the same mix of 0s through 9s. This is the easy, ideal scenario.

Non-IID means each client's data is skewed. Some clients might have almost all "3s" and "7s" while another gets mostly "0s" and "1s". This is what happens in the real world, and it makes federated learning much harder.

  • Open data.py in the Explorer sidebar and paste the imports and the dirichlet_partition() function:
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms


def dirichlet_partition(dataset, num_clients, alpha, seed=42):
    # Create a reproducible random number generator
    rng = np.random.default_rng(seed)
    targets = np.array(dataset.targets)
    num_classes = len(np.unique(targets))

    # Each client gets its own list of sample indices
    client_indices = [[] for _ in range(num_clients)]

    # For each digit class, split its samples across clients
    for c in range(num_classes):
        class_indices = np.where(targets == c)[0]
        rng.shuffle(class_indices)

        # Draw random proportions from the Dirichlet distribution
        proportions = rng.dirichlet(np.repeat(alpha, num_clients))
        proportions = proportions / proportions.sum()
        counts = (proportions * len(class_indices)).astype(int)

        # Distribute leftover samples from integer rounding
        remainder = len(class_indices) - counts.sum()
        for i in range(remainder):
            counts[i % num_clients] += 1

        # Assign this class's samples to each client by index
        start = 0
        for k in range(num_clients):
            end = start + counts[k]
            client_indices[k].extend(class_indices[start:end].tolist())
            start = end

    return client_indices

What does this code do?

  • matplotlib.use("Agg") tells matplotlib to render plots to files instead of opening a GUI window.
  • For each of the 10 digit classes, dirichlet_partition() draws a probability vector from Dir(alpha) with length equal to num_clients. These proportions decide how many samples of that class each client receives.
  • The remainder loop handles leftover samples from integer rounding so that no data is lost during partitioning.
  • Save data.py.

You will not see a change yet. The next function, get_client_loaders(), downloads the MNIST dataset and calls dirichlet_partition() to split it across clients.

  • Add the get_client_loaders() function below dirichlet_partition() in data.py:

def get_client_loaders(num_clients, alpha, batch_size=32, seed=42):
    # Standard MNIST preprocessing: convert to tensor and normalize
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])

    # Download MNIST if not already present in ./data
    train_dataset = datasets.MNIST(
        root="./data", train=True, download=True, transform=transform,
    )
    test_dataset = datasets.MNIST(
        root="./data", train=False, download=True, transform=transform,
    )

    # Partition training data across clients using Dirichlet
    client_indices = dirichlet_partition(train_dataset, num_clients, alpha, seed)

    # Create one DataLoader per client for batched training
    client_loaders = []
    for indices in client_indices:
        subset = Subset(train_dataset, indices)
        loader = DataLoader(subset, batch_size=batch_size, shuffle=True)
        client_loaders.append(loader)

    # Global test loader shared across all clients for evaluation
    test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

    return client_loaders, test_loader

What does this code do?

  • The normalization values 0.1307 (mean) and 0.3081 (std) are the dataset-wide statistics for MNIST. Normalizing inputs helps the model train faster.
  • A DataLoader is PyTorch's way of feeding data to a model in batches during training. Each client gets its own loader wrapping a Subset of the full MNIST dataset.
  • A global test loader covers the entire 10,000-sample test set. This is used later to evaluate the global model's accuracy after aggregation.
  • Save data.py.
  • Verify the data partitioning works by running this command in your terminal:
python -c "from data import get_client_loaders; cl, tl = get_client_loaders(10, 0.1); print([len(l.dataset) for l in cl])"
  • You should see a list of 10 numbers that vary significantly across clients (some may have 2x to 10x more samples than others). MNIST downloads automatically on the first run (about 60 MB).

MNIST download stalling or seeing a NameError?

  • If the download hangs, check your internet connection. The MNIST dataset is about 60 MB.
  • If you see a NameError for dirichlet_partition, make sure the function from the previous chunk is above get_client_loaders() in the same data.py file.
  • If you see a ModuleNotFoundError for torchvision, confirm your virtual environment is activated (look for (venv) in your prompt).

Why is my data partitioning failing?

  • Add the plot_data_distribution() function below get_client_loaders() in data.py:

def plot_data_distribution(client_indices, dataset, alpha, save_path="results"):
    os.makedirs(save_path, exist_ok=True)
    targets = np.array(dataset.targets)
    num_clients = len(client_indices)
    num_classes = 10

    # Build a matrix: rows = clients, columns = digit classes
    distribution = np.zeros((num_clients, num_classes))
    for k, indices in enumerate(client_indices):
        for idx in indices:
            distribution[k][targets[idx]] += 1

    # Render heatmap showing sample counts per client per class
    fig, ax = plt.subplots(figsize=(10, 6))
    im = ax.imshow(distribution, aspect="auto", cmap="YlOrRd")
    ax.set_xlabel("Digit Class")
    ax.set_ylabel("Client ID")
    ax.set_title(f"Data Distribution Across Clients (alpha={alpha})")
    ax.set_xticks(range(num_classes))
    ax.set_yticks(range(num_clients))
    plt.colorbar(im, ax=ax, label="Number of Samples")
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, f"data_distribution_alpha{alpha}.png"), dpi=150)
    plt.close()

What does this code do?

This function creates a heatmap showing how many samples of each digit class each client received. Rows are clients, columns are digit classes (0 through 9), and the color intensity shows the sample count.

When you see this heatmap for alpha=0.1, the non-IID problem becomes visually obvious. Some rows will have bright spots in only 1 or 2 columns, meaning those clients hold mostly a single digit class.

  • Save data.py.
  • Verify the complete file works by running:
python -c "from data import dirichlet_partition, get_client_loaders, plot_data_distribution; print('All functions imported successfully')"
  • You should see All functions imported successfully in your terminal. The heatmap will be generated later when you run the full experiment in a future step.

Seeing a SyntaxError or NameError?

  • Make sure all three functions (dirichlet_partition, get_client_loaders, plot_data_distribution) are in the same data.py file, in that order.
  • Check that the imports from the first chunk are still at the top of data.py. If you accidentally deleted them while pasting, re-add them.

Why is my data.py failing to import?

✔️ Awesome, I've got everything!

Great. Double check you have saved data.py with all three functions.

ⓧ I'd like to double check the full code

import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms


def dirichlet_partition(dataset, num_clients, alpha, seed=42):
    rng = np.random.default_rng(seed)
    targets = np.array(dataset.targets)
    num_classes = len(np.unique(targets))

    client_indices = [[] for _ in range(num_clients)]

    for c in range(num_classes):
        class_indices = np.where(targets == c)[0]
        rng.shuffle(class_indices)

        proportions = rng.dirichlet(np.repeat(alpha, num_clients))
        proportions = proportions / proportions.sum()
        counts = (proportions * len(class_indices)).astype(int)

        remainder = len(class_indices) - counts.sum()
        for i in range(remainder):
            counts[i % num_clients] += 1

        start = 0
        for k in range(num_clients):
            end = start + counts[k]
            client_indices[k].extend(class_indices[start:end].tolist())
            start = end

    return client_indices


def get_client_loaders(num_clients, alpha, batch_size=32, seed=42):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])

    train_dataset = datasets.MNIST(
        root="./data", train=True, download=True, transform=transform,
    )
    test_dataset = datasets.MNIST(
        root="./data", train=False, download=True, transform=transform,
    )

    client_indices = dirichlet_partition(train_dataset, num_clients, alpha, seed)

    client_loaders = []
    for indices in client_indices:
        subset = Subset(train_dataset, indices)
        loader = DataLoader(subset, batch_size=batch_size, shuffle=True)
        client_loaders.append(loader)

    test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

    return client_loaders, test_loader


def plot_data_distribution(client_indices, dataset, alpha, save_path="results"):
    os.makedirs(save_path, exist_ok=True)
    targets = np.array(dataset.targets)
    num_clients = len(client_indices)
    num_classes = 10

    distribution = np.zeros((num_clients, num_classes))
    for k, indices in enumerate(client_indices):
        for idx in indices:
            distribution[k][targets[idx]] += 1

    fig, ax = plt.subplots(figsize=(10, 6))
    im = ax.imshow(distribution, aspect="auto", cmap="YlOrRd")
    ax.set_xlabel("Digit Class")
    ax.set_ylabel("Client ID")
    ax.set_title(f"Data Distribution Across Clients (alpha={alpha})")
    ax.set_xticks(range(num_classes))
    ax.set_yticks(range(num_clients))
    plt.colorbar(im, ax=ax, label="Number of Samples")
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, f"data_distribution_alpha{alpha}.png"), dpi=150)
    plt.close()

Your model and data pipeline are ready. Next up, you write the client training loop and the FedAvg aggregation algorithm that ties everything together.

Implement Client Local Training

You have a CNN model and a data partitioner that splits MNIST across simulated clients. The next piece of the puzzle is what happens on each client's side.

In Federated Learning, every client (hospital) receives the current global model, trains it on its own private data, and sends back the updated weights. In this step, you will write the function that performs that local training using PyTorch's core training loop.

In this step, get ready to:
  • Write the client_update() function that trains a model copy on one client's local data.
  • Verify that local training returns updated model parameters and a sample count.
Write the client_update() function

Each client in FedAvg does the same thing every communication round. It receives the current global model, makes an independent copy, trains that copy on its own local data for a few epochs, and sends back the trained parameters.

Why deep-copy the model?

In Python, passing an object to a function does not create a new copy. If all 10 clients modified the same model object, each client's training would overwrite the previous client's work.

copy.deepcopy(model) creates a fully independent clone. Each client trains its own copy, and the original global model stays untouched.

  • Open client.py in the editor, then paste the following code:
import copy

import torch
import torch.nn as nn


def client_update(model, loader, lr=0.01, local_epochs=5, device="cpu"):
    # Create an independent copy so the global model stays untouched
    local_model = copy.deepcopy(model)
    local_model.to(device)
    local_model.train()

    # Set up SGD optimizer and cross-entropy loss
    optimizer = torch.optim.SGD(local_model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    # Train for the specified number of local epochs
    for _epoch in range(local_epochs):
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)

            # Core training loop: zero gradients, forward, loss, backward, update
            optimizer.zero_grad()
            outputs = local_model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

    # Return trained parameters and dataset size (used for weighted averaging)
    return local_model.state_dict(), len(loader.dataset)

What does this code do?

  • copy.deepcopy(model) creates a fully independent clone of the global model. This client trains its own copy without affecting anyone else.
  • local_model.train() puts the model in training mode. This enables behaviors like dropout that should only be active during training, not evaluation.
  • The inner loop (zero_grad, forward pass, loss.backward(), optimizer.step()) is the core PyTorch training pattern you will use in every future project.
  • The function returns two things: local_model.state_dict() (a dictionary mapping layer names to trained weight tensors) and len(loader.dataset) (the number of samples this client trained on). The server needs the sample count to give clients with more data proportionally more influence during aggregation.

What is a state_dict?

In PyTorch, a state_dict is a Python dictionary that maps each layer's name to its weight tensor. For your MNISTCnn, it contains entries like conv1.weight, fc1.bias, and so on.

This is exactly what gets "sent" from each client to the server. In a real deployment, you would serialize this dictionary and transmit it over the network. The server never sees the client's data.

  • Save client.py.
  • Verify the function works by running this command in your terminal:
python -c "import torch; from torch.utils.data import DataLoader, TensorDataset; from model import MNISTCnn; from client import client_update; loader = DataLoader(TensorDataset(torch.randn(8,1,28,28), torch.randint(0,10,(8,))), batch_size=4); sd, n = client_update(MNISTCnn(), loader, local_epochs=1); print(f'Returned {len(sd)} parameter tensors and {n} samples')"

You should see Returned 8 parameter tensors and 8 samples printed in your terminal. The 8 tensors correspond to the weight and bias for each of the four layers in your CNN (2 conv + 2 FC).

Seeing an import error?

  • Make sure you saved client.py before running the command. Unsaved changes are not picked up by Python.
  • Check that your terminal is in the fedavg-from-scratch directory, not a parent or child folder. Both model.py and client.py must be in the same directory for the imports to work.
  • Confirm your virtual environment is active. You should see (venv) at the start of your terminal prompt.

Help me debug my client_update() import error in PyTorch.

✔️ Awesome, I've got everything!

Great. Double check that you have saved client.py.

ⓧ I'd like to double check the full code

Here is the complete client.py file. Make sure yours matches exactly.

import copy

import torch
import torch.nn as nn


def client_update(model, loader, lr=0.01, local_epochs=5, device="cpu"):
    local_model = copy.deepcopy(model)
    local_model.to(device)
    local_model.train()

    optimizer = torch.optim.SGD(local_model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    for _epoch in range(local_epochs):
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = local_model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

    return local_model.state_dict(), len(loader.dataset)

Why PyTorch's training loop is explicit

If you have used TensorFlow/Keras before, you might be used to calling model.fit() and having everything happen automatically. PyTorch makes you write each step yourself: zero the gradients, run the forward pass, compute the loss, backpropagate, and update the weights.

This extra control is exactly why the federated learning research community uses PyTorch. When you need to modify the training loop (adding differential privacy noise, proximal regularization, or ZK proof hooks), you can do it line by line.

Your client-side training is ready. Each client can now receive a model, train it locally, and hand back the results. Next up, you will build the server that takes all those client updates and averages them into a single improved global model.

Implement Federated Averaging

You have the CNN model, the Dirichlet data partitioner, and the client training function ready. Now it is time to build the other half of federated learning: the server.

The server takes the trained model parameters from every client and combines them into a single improved global model. This is the FedAvg algorithm from McMahan et al. 2017, and the function your BOPFL research will eventually wrap with ZK proofs.

In this step, get ready to:
  • Write the FedAvg aggregation function that computes a weighted average of client model parameters.
  • Build the full federated training loop that orchestrates clients and server across communication rounds.
  • Add plotting and results output to compare IID and non-IID experiments.
Implement FedAvg server aggregation

The FedAvg formula is surprisingly simple. For each model parameter, take a weighted average across all clients, where the weight is the fraction of total data that client holds.

Why weighted averaging?

If Hospital A has 10,000 patient records and Hospital B has 100, you want Hospital A's trained model to count for more. It learned from more data.

The formula is: w_global = sum(n_k / n_total * w_k) for all participating clients.

  • Open server.py in your editor and add the FedAvg aggregation function:
import torch
import torch.nn as nn


def fedavg_aggregate(global_model, client_updates):
    # Calculate total samples across all clients for weighting
    total_samples = sum(n for _, n in client_updates)

    # Initialize aggregated state dict with zeros (same shape as model)
    global_state = global_model.state_dict()
    aggregated = {key: torch.zeros_like(param) for key, param in global_state.items()}

    # Weighted sum: each client's contribution proportional to its data size
    for state_dict, num_samples in client_updates:
        weight = num_samples / total_samples
        for key in aggregated:
            aggregated[key] += weight * state_dict[key].float()

    # Load averaged parameters into the global model
    global_model.load_state_dict(aggregated)
    return global_model

What does this code do?

  • total_samples sums the dataset size from every client. This becomes the denominator for the weight calculation.
  • aggregated starts as a dictionary of zero tensors matching the model's shape. Each client's parameters get multiplied by their weight (n_k / total_samples) and added to this running sum.
  • load_state_dict() replaces the global model's parameters with the freshly averaged values.

You won't see a change yet. The next chunk adds the evaluate() function to complete the server module.

The server also needs a way to measure how well the global model is doing after each round of aggregation.

  • Add the evaluate() function below fedavg_aggregate() in server.py:

def evaluate(model, test_loader, device="cpu"):
    model.to(device)
    # Switch to evaluation mode (disables dropout, batchnorm training behavior)
    model.eval()

    correct = 0
    total = 0
    total_loss = 0.0
    criterion = nn.CrossEntropyLoss()

    # Disable gradient computation since we are only measuring, not training
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * len(labels)
            _, predicted = torch.max(outputs, 1)
            total += len(labels)
            correct += (predicted == labels).sum().item()

    accuracy = correct / total
    avg_loss = total_loss / total
    return accuracy, avg_loss

What does this code do?

  • model.eval() switches the model to evaluation mode. Some layers (like dropout and batch normalization) behave differently during training vs. evaluation.
  • torch.no_grad() disables gradient computation. Evaluation does not need backpropagation, so this saves memory and speeds things up.
  • The function returns both accuracy (fraction of correct predictions) and avg_loss (average cross-entropy loss over the test set).
  • Save server.py.
  • Verify the aggregation function works by running this sanity check in your terminal:
python -c "from model import MNISTCnn; from server import fedavg_aggregate; m = MNISTCnn(); s = m.state_dict(); fedavg_aggregate(m, [(s, 100), (s, 100)]); print('Aggregation OK')"

You should see Aggregation OK printed to your terminal. This creates two identical state dicts, aggregates them with equal weights, and confirms the result loads without error.

Seeing an import error?

  • Make sure your terminal is in the fedavg-from-scratch folder. Python needs to find model.py and server.py in the current directory.
  • Check that your virtual environment is still active. You should see (venv) in your terminal prompt.
  • Confirm you saved server.py before running the command.

Why isn't my server.py importing correctly?

✔️ Awesome, I've got everything!

Great. Double check you have saved server.py.

ⓧ I'd like to double check the full code

import torch
import torch.nn as nn


def fedavg_aggregate(global_model, client_updates):
    total_samples = sum(n for _, n in client_updates)

    global_state = global_model.state_dict()
    aggregated = {key: torch.zeros_like(param) for key, param in global_state.items()}

    for state_dict, num_samples in client_updates:
        weight = num_samples / total_samples
        for key in aggregated:
            aggregated[key] += weight * state_dict[key].float()

    global_model.load_state_dict(aggregated)
    return global_model


def evaluate(model, test_loader, device="cpu"):
    model.to(device)
    model.eval()

    correct = 0
    total = 0
    total_loss = 0.0
    criterion = nn.CrossEntropyLoss()

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * len(labels)
            _, predicted = torch.max(outputs, 1)
            total += len(labels)
            correct += (predicted == labels).sum().item()

    accuracy = correct / total
    avg_loss = total_loss / total
    return accuracy, avg_loss
Implement the federated training loop

Now you need a function that orchestrates the full FedAvg protocol. Each communication round, the server sends the global model to every client, each client trains locally, and the server aggregates the results.

What is a communication round?

One round of FedAvg follows this cycle: (1) select clients, (2) send global model to each, (3) each client trains on its local data, (4) server collects updates and averages them, (5) evaluate the new global model.

The parameters you will use follow the paper's conventions: K=10 clients, T=50 rounds, E=5 local epochs, lr=0.01, B=32 batch size. All 10 clients participate every round (C=1.0).

  • Open train.py in your editor and add the imports and run_experiment() function:
import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch
from torchvision import datasets, transforms

from client import client_update
from data import dirichlet_partition, get_client_loaders, plot_data_distribution
from model import MNISTCnn
from server import evaluate, fedavg_aggregate


def run_experiment(
    num_clients=10,
    num_rounds=50,
    local_epochs=5,
    lr=0.01,
    batch_size=32,
    alpha=0.1,
    fraction_clients=1.0,
    seed=42,
    device="cpu",
):
    # Set seed for reproducibility
    torch.manual_seed(seed)

    # Create client data loaders with Dirichlet partitioning
    client_loaders, test_loader = get_client_loaders(
        num_clients, alpha, batch_size, seed,
    )

    print(f"\n{'=' * 60}")
    print(
        f"FedAvg Experiment: alpha={alpha}, K={num_clients}, "
        f"E={local_epochs}, rounds={num_rounds}"
    )
    print(f"{'=' * 60}")
    for i, loader in enumerate(client_loaders):
        print(f"  Client {i}: {len(loader.dataset)} samples")

    # Initialize global model
    global_model = MNISTCnn().to(device)

    # Track metrics across rounds
    history = {"round": [], "accuracy": [], "loss": []}

    # Evaluate the untrained model (round 0 baseline)
    acc, loss = evaluate(global_model, test_loader, device)
    history["round"].append(0)
    history["accuracy"].append(acc)
    history["loss"].append(loss)
    print(f"  Round 0 | Accuracy: {acc:.4f} | Loss: {loss:.4f}")

    num_selected = max(1, int(fraction_clients * num_clients))

    # Main FedAvg loop: select clients, train locally, aggregate
    for round_num in range(1, num_rounds + 1):
        if fraction_clients < 1.0:
            selected = torch.randperm(num_clients)[:num_selected].tolist()
        else:
            selected = list(range(num_clients))

        # Collect local updates from each selected client
        client_updates = []
        for k in selected:
            state_dict, n_samples = client_update(
                global_model, client_loaders[k], lr, local_epochs, device,
            )
            client_updates.append((state_dict, n_samples))

        # Aggregate updates using FedAvg
        global_model = fedavg_aggregate(global_model, client_updates)

        # Evaluate the updated global model
        acc, loss = evaluate(global_model, test_loader, device)
        history["round"].append(round_num)
        history["accuracy"].append(acc)
        history["loss"].append(loss)

        if round_num % 5 == 0 or round_num == 1:
            print(f"  Round {round_num} | Accuracy: {acc:.4f} | Loss: {loss:.4f}")

    return history

What does this code do?

  • The imports pull in every module you have built so far: MNISTCnn from model.py, the data functions from data.py, client_update from client.py, and both server functions.
  • matplotlib.use("Agg") tells matplotlib to render plots to files instead of opening a GUI window.
  • The main loop runs for num_rounds rounds. Each round calls client_update() for every selected client, then fedavg_aggregate() to combine them.
  • history records accuracy and loss after every round for later plotting.

Why does this all run in one Python process?

In a real FL deployment, each client would be a separate machine. But the algorithm is identical. This is how every FL researcher prototypes: simulate locally, deploy later.

  • Save train.py.

You won't see output yet. The __main__ block in the next substep calls this function.

Add plotting and results output

The training loop produces round-by-round metrics, but you need functions to visualize and save those results. The comparison plot showing IID vs non-IID convergence is the key deliverable for your README.

  • In train.py, add the plot_comparison() function below run_experiment():

def plot_comparison(iid_history, noniid_history, save_path="results"):
    os.makedirs(save_path, exist_ok=True)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    # Left plot: accuracy curves for both experiments
    ax1.plot(
        iid_history["round"], iid_history["accuracy"],
        label="IID (alpha=100)", linewidth=2,
    )
    ax1.plot(
        noniid_history["round"], noniid_history["accuracy"],
        label="Non-IID (alpha=0.1)", linewidth=2,
    )
    ax1.set_xlabel("Communication Round")
    ax1.set_ylabel("Test Accuracy")
    ax1.set_title("FedAvg: IID vs Non-IID Convergence")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # Right plot: loss curves for both experiments
    ax2.plot(
        iid_history["round"], iid_history["loss"],
        label="IID (alpha=100)", linewidth=2,
    )
    ax2.plot(
        noniid_history["round"], noniid_history["loss"],
        label="Non-IID (alpha=0.1)", linewidth=2,
    )
    ax2.set_xlabel("Communication Round")
    ax2.set_ylabel("Test Loss")
    ax2.set_title("FedAvg: IID vs Non-IID Loss")
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "fedavg_comparison.png"), dpi=150)
    plt.close()
    print(f"\nPlot saved to {save_path}/fedavg_comparison.png")

What does this code do?

This creates a side-by-side figure with accuracy vs. rounds on the left and loss vs. rounds on the right. Both IID and non-IID curves appear on each plot so you can see the convergence gap at a glance.

You won't see output yet. The __main__ block below calls this function.

Now add the results saver and the entry point that ties the whole pipeline together.

  • Add save_results() and the if __name__ == "__main__" block below plot_comparison() in train.py:

def save_results(iid_history, noniid_history, save_path="results"):
    # Save metrics as JSON for later analysis
    os.makedirs(save_path, exist_ok=True)
    results = {
        "iid_alpha_100": {
            "final_accuracy": iid_history["accuracy"][-1],
            "final_loss": iid_history["loss"][-1],
            "history": iid_history,
        },
        "noniid_alpha_0.1": {
            "final_accuracy": noniid_history["accuracy"][-1],
            "final_loss": noniid_history["loss"][-1],
            "history": noniid_history,
        },
    }
    filepath = os.path.join(save_path, "metrics.json")
    with open(filepath, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Metrics saved to {filepath}")


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    # Run IID experiment (balanced data across clients)
    iid_history = run_experiment(alpha=100, device=device)

    # Run non-IID experiment (skewed data across clients)
    noniid_history = run_experiment(alpha=0.1, device=device)

    # Generate comparison plots and save metrics
    plot_comparison(iid_history, noniid_history)
    save_results(iid_history, noniid_history)

    # Generate data distribution heatmaps
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    train_dataset = datasets.MNIST(
        root="./data", train=True, download=True, transform=transform,
    )
    noniid_indices = dirichlet_partition(train_dataset, 10, alpha=0.1, seed=42)
    plot_data_distribution(noniid_indices, train_dataset, alpha=0.1)
    iid_indices = dirichlet_partition(train_dataset, 10, alpha=100, seed=42)
    plot_data_distribution(iid_indices, train_dataset, alpha=100)

    # Print results table
    print(f"\n{'=' * 60}")
    print("RESULTS TABLE")
    print(f"{'=' * 60}")
    print(f"{'Setting':<25} {'Final Acc':>12} {'Final Loss':>12}")
    print(f"{'-' * 25} {'-' * 12} {'-' * 12}")
    print(
        f"{'IID (alpha=100)':<25} "
        f"{iid_history['accuracy'][-1]:>11.4f} "
        f"{iid_history['loss'][-1]:>11.4f}"
    )
    print(
        f"{'Non-IID (alpha=0.1)':<25} "
        f"{noniid_history['accuracy'][-1]:>11.4f} "
        f"{noniid_history['loss'][-1]:>11.4f}"
    )

What does the __main__ block do?

  • It runs two experiments back to back: one with alpha=100 (IID, balanced data) and one with alpha=0.1 (non-IID, skewed data).
  • After both experiments finish, it generates the side-by-side comparison plot, saves metrics to results/metrics.json, and creates data distribution heatmaps for both settings.
  • The final results table prints to the terminal so you can see the accuracy gap between IID and non-IID at a glance.
  • Save train.py.
  • Verify the full pipeline is wired together by running this quick import check:
python -c "from train import run_experiment; print('Pipeline ready')"

You should see Pipeline ready printed with no errors. This confirms all imports resolve and the function is callable.

Seeing a ModuleNotFoundError?

  • Check that server.py is saved with both fedavg_aggregate and evaluate. The train.py imports both of them.
  • Verify all five Python files (model.py, data.py, client.py, server.py, train.py) are in the same fedavg-from-scratch folder.
  • Make sure your terminal is in the fedavg-from-scratch directory and the (venv) prefix is visible.

Why can't Python find my modules when importing train.py?

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py.

ⓧ I'd like to double check the full code

import json
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch
from torchvision import datasets, transforms

from client import client_update
from data import dirichlet_partition, get_client_loaders, plot_data_distribution
from model import MNISTCnn
from server import evaluate, fedavg_aggregate


def run_experiment(
    num_clients=10,
    num_rounds=50,
    local_epochs=5,
    lr=0.01,
    batch_size=32,
    alpha=0.1,
    fraction_clients=1.0,
    seed=42,
    device="cpu",
):
    torch.manual_seed(seed)

    client_loaders, test_loader = get_client_loaders(
        num_clients, alpha, batch_size, seed,
    )

    print(f"\n{'=' * 60}")
    print(
        f"FedAvg Experiment: alpha={alpha}, K={num_clients}, "
        f"E={local_epochs}, rounds={num_rounds}"
    )
    print(f"{'=' * 60}")
    for i, loader in enumerate(client_loaders):
        print(f"  Client {i}: {len(loader.dataset)} samples")

    global_model = MNISTCnn().to(device)

    history = {"round": [], "accuracy": [], "loss": []}

    acc, loss = evaluate(global_model, test_loader, device)
    history["round"].append(0)
    history["accuracy"].append(acc)
    history["loss"].append(loss)
    print(f"  Round 0 | Accuracy: {acc:.4f} | Loss: {loss:.4f}")

    num_selected = max(1, int(fraction_clients * num_clients))

    for round_num in range(1, num_rounds + 1):
        if fraction_clients < 1.0:
            selected = torch.randperm(num_clients)[:num_selected].tolist()
        else:
            selected = list(range(num_clients))

        client_updates = []
        for k in selected:
            state_dict, n_samples = client_update(
                global_model, client_loaders[k], lr, local_epochs, device,
            )
            client_updates.append((state_dict, n_samples))

        global_model = fedavg_aggregate(global_model, client_updates)

        acc, loss = evaluate(global_model, test_loader, device)
        history["round"].append(round_num)
        history["accuracy"].append(acc)
        history["loss"].append(loss)

        if round_num % 5 == 0 or round_num == 1:
            print(f"  Round {round_num} | Accuracy: {acc:.4f} | Loss: {loss:.4f}")

    return history


def plot_comparison(iid_history, noniid_history, save_path="results"):
    os.makedirs(save_path, exist_ok=True)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    ax1.plot(
        iid_history["round"], iid_history["accuracy"],
        label="IID (alpha=100)", linewidth=2,
    )
    ax1.plot(
        noniid_history["round"], noniid_history["accuracy"],
        label="Non-IID (alpha=0.1)", linewidth=2,
    )
    ax1.set_xlabel("Communication Round")
    ax1.set_ylabel("Test Accuracy")
    ax1.set_title("FedAvg: IID vs Non-IID Convergence")
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    ax2.plot(
        iid_history["round"], iid_history["loss"],
        label="IID (alpha=100)", linewidth=2,
    )
    ax2.plot(
        noniid_history["round"], noniid_history["loss"],
        label="Non-IID (alpha=0.1)", linewidth=2,
    )
    ax2.set_xlabel("Communication Round")
    ax2.set_ylabel("Test Loss")
    ax2.set_title("FedAvg: IID vs Non-IID Loss")
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "fedavg_comparison.png"), dpi=150)
    plt.close()
    print(f"\nPlot saved to {save_path}/fedavg_comparison.png")


def save_results(iid_history, noniid_history, save_path="results"):
    os.makedirs(save_path, exist_ok=True)
    results = {
        "iid_alpha_100": {
            "final_accuracy": iid_history["accuracy"][-1],
            "final_loss": iid_history["loss"][-1],
            "history": iid_history,
        },
        "noniid_alpha_0.1": {
            "final_accuracy": noniid_history["accuracy"][-1],
            "final_loss": noniid_history["loss"][-1],
            "history": noniid_history,
        },
    }
    filepath = os.path.join(save_path, "metrics.json")
    with open(filepath, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Metrics saved to {filepath}")


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    iid_history = run_experiment(alpha=100, device=device)

    noniid_history = run_experiment(alpha=0.1, device=device)

    plot_comparison(iid_history, noniid_history)
    save_results(iid_history, noniid_history)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    train_dataset = datasets.MNIST(
        root="./data", train=True, download=True, transform=transform,
    )
    noniid_indices = dirichlet_partition(train_dataset, 10, alpha=0.1, seed=42)
    plot_data_distribution(noniid_indices, train_dataset, alpha=0.1)
    iid_indices = dirichlet_partition(train_dataset, 10, alpha=100, seed=42)
    plot_data_distribution(iid_indices, train_dataset, alpha=100)

    print(f"\n{'=' * 60}")
    print("RESULTS TABLE")
    print(f"{'=' * 60}")
    print(f"{'Setting':<25} {'Final Acc':>12} {'Final Loss':>12}")
    print(f"{'-' * 25} {'-' * 12} {'-' * 12}")
    print(
        f"{'IID (alpha=100)':<25} "
        f"{iid_history['accuracy'][-1]:>11.4f} "
        f"{iid_history['loss'][-1]:>11.4f}"
    )
    print(
        f"{'Non-IID (alpha=0.1)':<25} "
        f"{noniid_history['accuracy'][-1]:>11.4f} "
        f"{noniid_history['loss'][-1]:>11.4f}"
    )

Your full FedAvg pipeline is wired together and ready to run. Next up, you will execute the experiments, analyze the IID vs non-IID convergence gap, and write tests that prove your implementation is correct.

Run Experiments, Analyze Results, and Write Tests

You have a complete FedAvg implementation spread across five files. Every piece is in place: the CNN model, the Dirichlet partitioner, the client training function, the server aggregation logic, and the experiment orchestrator.

Now it is time to stop coding and start thinking like a researcher. You will run the full IID vs non-IID experiment, interpret the convergence gap, prove your code is correct with pytest, and write everything up in a research-audience README.

In this step, get ready to:
  • Run 50-round IID and non-IID experiments and compare convergence curves.
  • Write a pytest suite that asserts model architecture, data partitioning, and FedAvg convergence.
  • Create a research-audience README and commit your project to Git.
Execute the full experiment

Your train.py script runs two complete FedAvg experiments back to back. The first uses alpha=100 (IID data, evenly distributed). The second uses alpha=0.1 (non-IID data, unevenly distributed).

  • In your terminal, run the experiment by executing this command:
python train.py

How long will this take?

On CPU, expect approximately 15-25 minutes total for both experiments (50 rounds each). You will see round-by-round accuracy printing to the terminal as training progresses.

Do not close the terminal or interrupt the process. The script handles everything automatically, including generating plots and saving metrics.

Once the script finishes, you should see a results table printed in the terminal showing final accuracy and loss for both settings.

  • Confirm the output files were generated by running:
dir results
  • You should see four files: fedavg_comparison.png, data_distribution_alpha0.1.png, data_distribution_alpha100.png, and metrics.json.
  • Open fedavg_comparison.png from the results folder in the Explorer sidebar to view your convergence curves.

What do the results mean?

  • IID (alpha=100): Final accuracy around 98-99%, with a smooth convergence curve. Every client had a balanced mix of all digits, so local training aligned well with the global objective.
  • Non-IID (alpha=0.1): Final accuracy around 90-95%, with slower and noisier convergence. Some clients held mostly one or two digit classes, so their local training pulled the model in conflicting directions.
  • The gap between these two curves is called client drift. This accuracy degradation under heterogeneous data is the core problem that motivates research into algorithms like FedProx.

What about the heatmaps?

Open data_distribution_alpha0.1.png. You will see bright spots concentrated in just 1-2 columns per row. This means some clients received mostly one digit class. Compare it with data_distribution_alpha100.png, where every row has a roughly even spread. This visual contrast is the non-IID problem in one image.

Not seeing accuracy improve?

  • Make sure your virtual environment is activated. You should see (venv) at the start of your terminal prompt.
  • If you see an import error, check that you are running the command from inside the fedavg-from-scratch directory where all your .py files live.
  • If accuracy stays near 10% (random chance for 10 classes), double check that client.py calls loss.backward() and optimizer.step() in the training loop.

Why is my FedAvg experiment not converging or showing import errors?

Write pytest tests

Testing ML code is a skill most beginners skip, but it is expected in research. Your tests verify that the implementation is correct, not just that it "seems to work."

You will write three test classes covering the model architecture, the data partitioner, and FedAvg convergence.

  • Open tests/test_fedavg.py in the Explorer sidebar.
  • Add the imports and the first test class by pasting the following code:
import torch
import pytest

from model import MNISTCnn
from data import dirichlet_partition, get_client_loaders
from client import client_update
from server import fedavg_aggregate, evaluate


class TestModel:
    def test_output_shape(self):
        # Feed random images through the CNN and check output dimensions
        model = MNISTCnn()
        x = torch.randn(4, 1, 28, 28)
        out = model(x)
        assert out.shape == (4, 10), f"Expected (4, 10), got {out.shape}"

    def test_parameter_count(self):
        # Verify total params match the McMahan et al. 2017 architecture
        model = MNISTCnn()
        total = sum(p.numel() for p in model.parameters())
        assert total == 1_663_370, f"Expected 1,663,370 params, got {total}"

What do these tests check?

  • test_output_shape feeds a batch of 4 random 28x28 images through MNISTCnn and asserts the output shape is (4, 10). This verifies all the layers are wired correctly from input to output.
  • test_parameter_count asserts the total parameter count equals 1,663,370. This confirms you matched the exact architecture from the McMahan et al. 2017 paper.
  • You will not see a change yet. The next chunk adds the data partitioning tests.
  • Add the TestDirichletPartition class below the TestModel class you just pasted:

class TestDirichletPartition:
    def test_all_samples_assigned(self):
        # Verify no data is lost during Dirichlet partitioning
        from torchvision import datasets, transforms

        transform = transforms.Compose([transforms.ToTensor()])
        dataset = datasets.MNIST(
            root="./data", train=True, download=True, transform=transform,
        )
        indices = dirichlet_partition(dataset, num_clients=10, alpha=0.1)
        total = sum(len(idx) for idx in indices)
        assert total == len(dataset), (
            f"Expected {len(dataset)} total samples, got {total}"
        )

    def test_noniid_creates_imbalance(self):
        # Verify alpha=0.1 actually produces uneven splits
        from torchvision import datasets, transforms

        transform = transforms.Compose([transforms.ToTensor()])
        dataset = datasets.MNIST(
            root="./data", train=True, download=True, transform=transform,
        )
        indices = dirichlet_partition(dataset, num_clients=10, alpha=0.1)
        sizes = [len(idx) for idx in indices]
        assert max(sizes) / min(sizes) > 2.0, (
            "Non-IID partition should create significant imbalance"
        )

What do these tests check?

  • test_all_samples_assigned partitions MNIST with alpha=0.1 and asserts the total samples across all clients equals 60,000. If any data was lost or duplicated during partitioning, this test catches it.
  • test_noniid_creates_imbalance checks that the ratio of the largest to smallest client dataset exceeds 2.0. This verifies the Dirichlet distribution actually creates uneven splits, not near-equal ones.
  • You will not see a change yet. The next chunk adds the FedAvg convergence tests.
  • Add the TestFedAvg class below the TestDirichletPartition class:

class TestFedAvg:
    def test_aggregation_preserves_shape(self):
        # Aggregating identical state dicts should preserve tensor shapes
        model = MNISTCnn()
        state1 = model.state_dict()
        state2 = model.state_dict()
        updates = [(state1, 100), (state2, 200)]
        result = fedavg_aggregate(model, updates)
        for key in state1:
            assert result.state_dict()[key].shape == state1[key].shape

    def test_convergence_iid(self):
        """FedAvg on IID MNIST must reach >90% accuracy in 10 rounds."""
        # The key correctness test: FedAvg must converge on IID data
        client_loaders, test_loader = get_client_loaders(
            num_clients=5, alpha=100.0, batch_size=64, seed=42,
        )
        model = MNISTCnn()
        for _round in range(10):
            updates = []
            for loader in client_loaders:
                sd, n = client_update(model, loader, lr=0.01, local_epochs=2)
                updates.append((sd, n))
            model = fedavg_aggregate(model, updates)
        acc, _ = evaluate(model, test_loader)
        assert acc > 0.90, f"Expected >90% accuracy after 10 rounds, got {acc:.4f}"

What do these tests check?

  • test_aggregation_preserves_shape aggregates two identical state dicts with different sample counts and asserts the output tensor shapes match. This is a sanity check that the aggregation function does not corrupt parameter dimensions.
  • test_convergence_iid is the most important test. It runs 10 rounds of FedAvg on 5 clients with IID data and asserts accuracy exceeds 90%. If FedAvg is implemented correctly, it must converge on IID data. A failure here means something in your training loop or aggregation is broken.
  • Save tests/test_fedavg.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved tests/test_fedavg.py before running the tests.

ⓧ I'd like to double check the full code

Here is the complete tests/test_fedavg.py file. Make sure yours matches exactly.

import torch
import pytest

from model import MNISTCnn
from data import dirichlet_partition, get_client_loaders
from client import client_update
from server import fedavg_aggregate, evaluate


class TestModel:
    def test_output_shape(self):
        model = MNISTCnn()
        x = torch.randn(4, 1, 28, 28)
        out = model(x)
        assert out.shape == (4, 10), f"Expected (4, 10), got {out.shape}"

    def test_parameter_count(self):
        model = MNISTCnn()
        total = sum(p.numel() for p in model.parameters())
        assert total == 1_663_370, f"Expected 1,663,370 params, got {total}"


class TestDirichletPartition:
    def test_all_samples_assigned(self):
        from torchvision import datasets, transforms

        transform = transforms.Compose([transforms.ToTensor()])
        dataset = datasets.MNIST(
            root="./data", train=True, download=True, transform=transform,
        )
        indices = dirichlet_partition(dataset, num_clients=10, alpha=0.1)
        total = sum(len(idx) for idx in indices)
        assert total == len(dataset), (
            f"Expected {len(dataset)} total samples, got {total}"
        )

    def test_noniid_creates_imbalance(self):
        from torchvision import datasets, transforms

        transform = transforms.Compose([transforms.ToTensor()])
        dataset = datasets.MNIST(
            root="./data", train=True, download=True, transform=transform,
        )
        indices = dirichlet_partition(dataset, num_clients=10, alpha=0.1)
        sizes = [len(idx) for idx in indices]
        assert max(sizes) / min(sizes) > 2.0, (
            "Non-IID partition should create significant imbalance"
        )


class TestFedAvg:
    def test_aggregation_preserves_shape(self):
        model = MNISTCnn()
        state1 = model.state_dict()
        state2 = model.state_dict()
        updates = [(state1, 100), (state2, 200)]
        result = fedavg_aggregate(model, updates)
        for key in state1:
            assert result.state_dict()[key].shape == state1[key].shape

    def test_convergence_iid(self):
        """FedAvg on IID MNIST must reach >90% accuracy in 10 rounds."""
        client_loaders, test_loader = get_client_loaders(
            num_clients=5, alpha=100.0, batch_size=64, seed=42,
        )
        model = MNISTCnn()
        for _round in range(10):
            updates = []
            for loader in client_loaders:
                sd, n = client_update(model, loader, lr=0.01, local_epochs=2)
                updates.append((sd, n))
            model = fedavg_aggregate(model, updates)
        acc, _ = evaluate(model, test_loader)
        assert acc > 0.90, f"Expected >90% accuracy after 10 rounds, got {acc:.4f}"
  • Run the test suite by executing this command:
pytest tests/test_fedavg.py -v

How long does this take?

The test_convergence_iid test actually trains a model for 10 rounds, so expect 2-3 minutes for the full suite. The other five tests finish in seconds.

  • You should see all 6 tests passing with green checkmarks. Every line should show PASSED.

Seeing test failures?

  • If test_parameter_count fails, check that your model.py uses the exact kernel sizes and layer dimensions from earlier (5x5 kernels, 32 and 64 filters, 512-unit FC layer).
  • If test_convergence_iid fails with accuracy below 90%, this usually means your client_update() or fedavg_aggregate() has a bug. Double check the training loop in client.py and the weighted averaging in server.py.
  • If you see ModuleNotFoundError, make sure you are running pytest from the fedavg-from-scratch directory, not from inside the tests folder.

Why are my pytest tests failing for the FedAvg project?

Write the README and commit to Git

The last deliverable is your README.md. This is your research writing practice. You will describe the problem, show the architecture, present your results, and reflect on what you learned.

  • Create a new file called README.md in the fedavg-from-scratch folder by right-clicking in the Explorer sidebar and selecting New File.
  • Paste the following template into README.md and fill in your own content where indicated:
# FedAvg from Scratch

## Problem Statement

<!-- Write 1-2 paragraphs explaining what FedAvg solves (training ML models
     across distributed data holders without centralizing data) and why
     non-IID data is the core challenge. Write as if explaining to a
     professor who knows ML but not FL. -->

## Architecture Diagram

```
                  +-------------------+
                  |      Server       |
                  |  (Global Model)   |
                  +---------+---------+
                            |
              Send global parameters to all clients
                            |
         +------------------+------------------+
         |                  |                  |
    +----v-----+      +----v-----+       +----v-----+
    | Client 0 |      | Client 1 |  ...  | Client 9 |
    | Local SGD|      | Local SGD|       | Local SGD|
    +----+-----+      +----+-----+       +----+-----+
         |                  |                  |
         +------------------+------------------+
                            |
              Send updated parameters to server
                            |
                  +---------v---------+
                  |  FedAvg Aggregate  |
                  |  w = sum(nk/N)*wk  |
                  +-------------------+
```

## Results Table

| Setting | Alpha | Final Accuracy | Final Loss |
|---------|-------|----------------|------------|
| IID     | 100   | <!-- your value --> | <!-- your value --> |
| Non-IID | 0.1   | <!-- your value --> | <!-- your value --> |

<!-- Fill in the actual numbers from your terminal output. -->

## What I Learned

<!-- Write 2-3 paragraphs for a research audience covering:
     1. The mechanical understanding you gained of FedAvg (e.g. weighted
        averaging formula, state_dict, the training loop pattern).
     2. The empirical impact of non-IID data on convergence (reference
        your plot, explain client drift).
     3. Connection to your broader research (why this degradation
        motivates algorithms like FedProx and DP). -->

What goes in each section?

  • Problem Statement: Explain what federated learning solves (distributed training without centralizing data) and why non-IID data is the core challenge. Write for someone who knows ML but not FL.
  • Architecture Diagram: The ASCII diagram is already provided in the template. It shows the server-client flow: global parameters go down, updated parameters come back up, and the server aggregates.
  • Results Table: Replace the placeholders with the actual numbers from the results table your terminal printed after running the experiment.
  • What I Learned: Write 2-3 paragraphs reflecting on the FedAvg mechanics you implemented, the convergence gap you observed, and why this matters for research into better aggregation algorithms.
  • Save README.md after filling in your content.

With all the code, tests, results, and documentation in place, it is time to commit everything to Git.

  • Stage and commit all your work by running these commands:
git add -A
git commit -m "feat: FedAvg from scratch with IID vs non-IID comparison"

What does this commit include?

The git add -A command stages all new and modified files. Your .gitignore from earlier automatically excludes venv/, data/, results/, and __pycache__/ from the commit. Only your source code, tests, and README get versioned.

  • Verify the commit was created by running:
git log --oneline
  • You should see your commit message feat: FedAvg from scratch with IID vs non-IID comparison at the top of the log.

Git commit not working?

  • If Git asks for your name and email, set them by running git config user.name "Your Name" and git config user.email "you@example.com", then retry the commit.
  • If you see a message about nothing to commit, make sure you saved all your files before running git add -A.

Why is my Git commit failing in this project?

Secret mission

Implement Krum Byzantine-Resilient Aggregation

Everything so far assumed all 10 clients are honest. But what if a malicious actor joins your federation and sends garbage model updates to sabotage training? In this mission, you will implement the Krum algorithm from Blanchard et al. 2017 to defend against exactly this attack. You will simulate a Byzantine client injecting random noise and prove that Krum resists it while FedAvg collapses.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • data/ directory containing the downloaded MNIST dataset (~60 MB).
  • venv/ directory containing your Python virtual environment and installed packages (PyTorch, matplotlib, pytest).

Keep everything running

No action needed. Choose this if you are still actively working on this project or plan to continue to the next project in the series.

  • The MNIST data in data/ and the virtual environment in venv/ are reused by Projects 2-4 in this series.
  • Everything runs locally on your machine. There are no cloud resources or subscriptions to worry about.

Pause - I'll come back to this later

Deactivate the virtual environment to free up your terminal. Everything persists on disk for when you return.

  • Deactivate your virtual environment by running:
deactivate

Your project files, MNIST data, and generated results all remain on disk. When you return, navigate back into your fedavg-from-scratch folder and reactivate the virtual environment.

macOS

  • Reactivate by running:
source venv/bin/activate

Windows

  • Reactivate by running:
venv\Scripts\activate

Delete - I don't want to use this again

Remove all project resources and start fresh.

  • Deactivate your virtual environment first by running:
deactivate

Keep data/ if you are continuing the series

The data/ folder contains the downloaded MNIST dataset (~60 MB). If you plan to continue to the next project in this series, keep it. The MNIST data is reused and saves you a re-download.

  • To delete just the virtual environment, remove the venv/ folder.

macOS

rm -rf venv

Windows

rmdir /s /q venv
  • To remove everything including the MNIST data and generated results, navigate out of the project folder and delete it entirely.

macOS

cd ~/Desktop
rm -rf fedavg-from-scratch

Windows

cd ~\Desktop
rmdir /s /q fedavg-from-scratch

Nice Work!

Nice Work!

Nice work! You just built a complete Federated Averaging implementation from scratch in PyTorch, with every line of the training loop and averaging math written and understood by you.

You've learned how to:

  • Built a 1,663,370-parameter CNN matching the McMahan et al. 2017 paper, learning the nn.Module, state_dict, and forward method pattern that replaces Keras's sequential API.
  • Implemented Dirichlet data partitioning and the complete FedAvg protocol from scratch, including local SGD on each client, weighted parameter averaging on the server, and 50-round communication loops.
  • Ran IID vs non-IID experiments producing convergence plots and a pytest suite with a convergence assertion, all documented in a research-audience README.md.
  • Secret Mission: Implemented Krum Byzantine-resilient aggregation and demonstrated that a single malicious client collapses FedAvg while Krum maintains convergence.

Ready to quiz yourself?