Federated Learning with Top-K Sparsification

Simulate IoT federated learning and cut communication costs by 90%.

Introduction

30 Second Summary

IoT devices like smart sensors and wearables generate valuable data, but they run on tiny batteries and slow networks. Training machine learning models across thousands of these devices means sending model updates back and forth constantly. That communication cost can drain a device's battery in hours.

In this project, you will build a federated learning simulation that trains a neural network across 10 IoT clients, measure exactly how many bytes each client transmits per round, and implement Top-K gradient sparsification to slash communication by up to 95%.

What You'll Build

You will run a federation of 10 simulated IoT clients that collaboratively train a CNN on image data, watch accuracy climb round by round, then flip on gradient compression and see the bytes-transmitted counter plummet. At the end, you produce a chart that reveals the exact sweet spot between saving bandwidth and preserving accuracy.

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

  • A 10-client federated simulation with realistic non-IID data splits using Dirichlet partitioning, where each client sees a different slice of CIFAR-10 images.
  • Per-round communication metering that tracks exactly how many bytes each client transmits to the server, exposing the full cost of uncompressed model updates.
  • A Top-K sparsification module that keeps only the largest model-update values and zeros out the rest, cutting transmitted data by up to 95% in a single line of config.
  • A comparison chart plotting final accuracy against total bytes transmitted across four sparsity levels, revealing the Pareto frontier you would present in a research paper.
  • Secret Mission: Add error feedback with residual accumulation to recover accuracy lost from aggressive sparsification.

Are there any prerequisites?

You should be comfortable with basic Python (functions, loops, dictionaries). Some familiarity with PyTorch (tensors, training loops) is helpful but not required. Every code block includes comments that explain what each piece does.

Before We Start

Before diving into code, take a moment to lock in what you are building and why it matters.

Set Up Your FL Development Environment

Before you can run any federated learning experiments, you need the right tools in place. Flower is the framework that orchestrates the federation, and PyTorch provides the neural network training engine underneath.

In this step, you will install both frameworks (along with supporting libraries) inside an isolated virtual environment, then scaffold the project files that Flower needs to run a simulation.

In this step, get ready to:
  • Create and activate a Python virtual environment.
  • Install Flower, PyTorch, and all supporting libraries.
  • Scaffold the project directory and Flower configuration file.
Create a virtual environment and activate it

A virtual environment keeps your project's packages separate from the rest of your system. This prevents version conflicts between different Python projects.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Terminal (macOS) or Command Prompt (Windows) and press Enter to open it.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir fl-topk
cd fl-topk
  • Create the virtual environment by running this command:
python -m venv fl-topk-env

What does this command do?

python -m venv fl-topk-env creates a self-contained Python installation inside a folder called fl-topk-env. Any packages you install while this environment is active stay inside it, keeping your system Python clean.

  • Activate the virtual environment by running this command:

macOS

source fl-topk-env/bin/activate

Windows

fl-topk-env\Scripts\activate

Your terminal prompt now shows (fl-topk-env) at the beginning. This confirms the virtual environment is active.

Don't see (fl-topk-env) in your prompt?

Make sure you ran the activation command for your operating system. On macOS/Linux, the command starts with source. On Windows, it uses a backslash path with no source prefix.

help me activate my Python virtual environment

Install all dependencies in one command

With the virtual environment active, every package you install goes into fl-topk-env instead of your system Python.

  • Install Flower, PyTorch, and all supporting libraries by running this command:
pip install "flwr[simulation]" torch torchvision "flwr-datasets[vision]" matplotlib

What is this installing?

  • flwr[simulation] installs Flower with its simulation runtime, so you can run a full federation on a single machine.
  • torch and torchvision provide PyTorch and its image-processing utilities for training the CNN.
  • flwr-datasets[vision] provides built-in dataset partitioning, including the DirichletPartitioner you will use for non-IID data splits.
  • matplotlib is the plotting library you will use to chart accuracy vs. communication cost.

This install may take a few minutes depending on your internet speed. You will see a stream of download and install progress lines in the terminal.

  • Once the install finishes, confirm Flower is installed by running:
flwr --version

You should see output like Flower version: 1.33.0 (or a higher number).

Seeing 'command not found'?

Make sure your virtual environment is still active. You should see (fl-topk-env) in your terminal prompt. If it disappeared, reactivate the environment using the activation command for your OS, then try flwr --version again.

help me fix flwr command not found after pip install

Create the project directory structure

Flower projects follow a specific layout. You need a Python package folder for your code and a pyproject.toml file that tells Flower where to find your client and server components.

  • Create the fl_topk package folder and its empty Python files by running these commands:

macOS

mkdir fl_topk
touch fl_topk/__init__.py fl_topk/task.py fl_topk/client_app.py fl_topk/server_app.py

Windows

mkdir fl_topk
type nul > fl_topk\__init__.py
type nul > fl_topk\task.py
type nul > fl_topk\client_app.py
type nul > fl_topk\server_app.py
  • Confirm the files exist by running:
ls fl_topk

You should see __init__.py, client_app.py, server_app.py, and task.py in the output.

What are these files for?

  • task.py will hold the neural network model, data loading, and utility functions.
  • client_app.py will define how each simulated IoT client trains and evaluates the model.
  • server_app.py will orchestrate the federation, aggregating updates from all clients each round.
  • __init__.py marks fl_topk as a Python package so Flower can import from it.

Now you need a code editor to create the pyproject.toml configuration file.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type VS Code and press Enter to open it. If you do not have VS Code installed, open any code editor you prefer.
  • In your editor, click File in the top menu bar.
  • Click Open Folder.
  • Select the fl-topk folder on your Desktop, then click Open.
  • Create a new file in the root of fl-topk (not inside the fl_topk subfolder) and name it pyproject.toml.

Creating a new file in VS Code

In the Explorer sidebar, hover over the FL-TOPK header. Click the New File icon (a page with a plus sign) and type pyproject.toml.

  • Paste the following project metadata into pyproject.toml:
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[project]
name = "fl-topk"
version = "1.0.0"
description = "Communication-efficient FL with Top-K gradient sparsification"
license = "Apache-2.0"
dependencies = [
    "flwr[simulation]>=1.33.0",
    "flwr-datasets[vision]>=0.6.0",
    "torch>=2.0.0",
    "torchvision>=0.15.0",
    "matplotlib>=3.7.0",
]

What does this section configure?

The [build-system] section tells Python to use Hatchling as the build backend. The [project] section names the project and declares its dependencies with minimum versions.

  • Below the code you just pasted, add the Flower configuration section:
[tool.flwr.app]
publisher = "nextwork"

[tool.flwr.app.components]
serverapp = "fl_topk.server_app:app"
clientapp = "fl_topk.client_app:app"

[tool.flwr.app.config]
num-server-rounds = 5
local-epochs = 2
lr = 0.01
sparsity-ratio = 0.0

[tool.flwr.federations]
default = "local-simulation"

[tool.flwr.federations.local-simulation]
options.num-supernodes = 10
options.backend.client-resources.num-cpus = 2

What does this section configure?

The [tool.flwr.app.components] section tells Flower where to find your server and client code. The [tool.flwr.app.config] section sets default hyperparameters: 5 training rounds, 2 local epochs per client, a learning rate of 0.01, and no sparsification yet.

The [tool.flwr.federations.local-simulation] section configures 10 simulated SuperNodes (one per IoT client) each with 2 CPU cores. This lets you run the entire federation locally on a single machine.

  • Save pyproject.toml.

✔️ Awesome, I've got everything!

Great. Double check you have saved pyproject.toml in the root of your fl-topk folder (not inside fl_topk).

ⓧ I'd like to double check the full code

Your pyproject.toml should contain two sections that you pasted above. Scroll up and confirm both chunks are present in this order:

  • The [build-system] and [project] sections (project metadata and dependencies).
  • The [tool.flwr.app] through [tool.flwr.federations.local-simulation] sections (Flower configuration).

Make sure there is a blank line between each section header and that the file is saved.

Not sure if pyproject.toml is in the right place?

The pyproject.toml file must be in the fl-topk root folder, at the same level as the fl_topk subfolder. If it is inside fl_topk, Flower will not find it. Check the Explorer sidebar in your editor to see the file hierarchy.

help me fix my pyproject.toml file location for a Flower project

Before you move on, predict what you think flwr --version will show now that everything is installed.

  • Switch back to your terminal and confirm the full setup by running:
flwr --version

You should see Flower version: 1.33.0 (or higher) printed. This confirms Flower, its simulation runtime, and all dependencies are correctly installed.

Your environment is ready and the project skeleton is in place. Next up, you will fill in these empty files to build and run a working federated learning simulation.

Build and Run a Baseline FedAvg Federation

Your development environment is set up with Flower, PyTorch, and all supporting libraries installed. Before you can optimize communication, you need a working federation to establish a baseline accuracy.

In this step, you will build a CNN model, partition CIFAR-10 data across 10 simulated IoT clients using Dirichlet partitioning, and run your first FedAvg experiment. By the end, you will watch accuracy climb across rounds in your terminal.

In this step, get ready to:
  • Define a CNN model and non-IID data loader for CIFAR-10 training.
  • Build Flower client and server apps that coordinate federated rounds.
  • Run the federation and observe accuracy improving over 5 rounds.
Define the model and data loading

The task.py file is the shared toolkit for both the client and the server. It holds the neural network architecture, data loading logic, and the training and testing functions that every federated participant uses.

  • Open fl_topk/task.py in your code editor and add the imports and CNN class by pasting the following code:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from flwr_datasets import FederatedDataset
from flwr_datasets.partitioner import DirichletPartitioner
from torchvision.transforms import Compose, Normalize, ToTensor

# Small CNN for CIFAR-10 (3 input channels, 10 output classes)
class Net(nn.Module):
    """Small CNN for CIFAR-10 classification."""

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # Two conv+pool stages, then flatten and classify
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 32 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

What does this code do?

The imports pull in PyTorch for the model, Flower Datasets for partitioned data loading, and torchvision transforms for image normalization.

The Net class is a small CNN designed for 32x32 color images. Two convolutional layers extract features, two max-pool layers halve the spatial dimensions, and two fully connected layers produce 10 class probabilities. The model has roughly 62,500 parameters. That is small enough to train on CPU while still producing meaningful communication costs to measure.

  • Save task.py.
  • Verify the model loads correctly by running this command in your terminal:
python -c "from fl_topk.task import Net; m = Net(); print(sum(p.numel() for p in m.parameters()), 'parameters')"

You should see 62506 parameters printed.

Seeing a ModuleNotFoundError?

  • Make sure you are running the command from the parent directory that contains the fl_topk folder, not from inside fl_topk itself.
  • Confirm your virtual environment is active. You should see (fl-topk-env) in your terminal prompt.
  • Check that fl_topk/__init__.py exists (even if empty). Python needs it to treat the folder as a package.

Still stuck? Why can't Python find my fl_topk module?

  • Below the Net class, add the data loading function that creates non-IID partitions:
fds = None

def load_data(partition_id: int, num_partitions: int):
    """Load a CIFAR-10 partition for a single client."""
    global fds
    if fds is None:
        # Create non-IID splits using Dirichlet distribution
        partitioner = DirichletPartitioner(
            num_partitions=num_partitions, partition_by="label", alpha=0.5
        )
        fds = FederatedDataset(
            dataset="uoft-cs/cifar10",
            partitioners={"train": partitioner},
        )
    partition = fds.load_partition(partition_id)
    # Normalize to standard CIFAR-10 mean and std values
    pytorch_transforms = Compose(
        [ToTensor(), Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))]
    )
    def apply_transforms(batch):
        batch["img"] = [pytorch_transforms(img) for img in batch["img"]]
        return batch
    partition = partition.with_transform(apply_transforms)
    trainloader = DataLoader(partition, batch_size=32, shuffle=True)
    # Load shared test set for evaluation
    testset = fds.load_split("test")
    testset = testset.with_transform(apply_transforms)
    testloader = DataLoader(testset, batch_size=64)
    return trainloader, testloader

What is Dirichlet partitioning?

In real IoT deployments, each device collects different kinds of data. A temperature sensor in a warehouse sees different patterns than one in a greenhouse. This is called non-IID (non-identically distributed) data.

DirichletPartitioner with alpha=0.5 simulates this by unevenly distributing CIFAR-10 labels across 10 clients. Some clients see mostly cats and dogs, others mostly trucks and planes. This is the standard approach used in FL research papers like FedProx and FedAdaptQ.

  • Below load_data, add the train function that each client calls to update the model on its local data:
def train(model, trainloader, epochs, lr, device):
    """Train the model on local data."""
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    model.train()
    total_loss = 0.0
    num_batches = 0
    for _ in range(epochs):
        for batch in trainloader:
            images = batch["img"].to(device)
            labels = batch["label"].to(device)
            optimizer.zero_grad()
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            num_batches += 1
    # Return average loss across all batches
    return total_loss / max(num_batches, 1)

What does this code do?

This is a standard PyTorch training loop. Each client calls this function to train the global model on its own local data for a configurable number of epochs. SGD with momentum 0.9 is the standard optimizer used in FL benchmarks.

  • Add the test function below train to evaluate the model without updating weights:
def test(model, testloader, device):
    """Evaluate the model."""
    criterion = nn.CrossEntropyLoss()
    model.eval()
    correct = 0
    total = 0
    total_loss = 0.0
    # Disable gradient computation for evaluation
    with torch.no_grad():
        for batch in testloader:
            images = batch["img"].to(device)
            labels = batch["label"].to(device)
            outputs = model(images)
            total_loss += criterion(outputs, labels).item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    accuracy = correct / total
    avg_loss = total_loss / len(testloader)
    return avg_loss, accuracy

What does this code do?

The test function evaluates the model on a held-out test set. torch.no_grad() disables gradient computation so no weights are updated. It returns both the average loss and accuracy, which the server will log after each round.

  • Save task.py.
  • Confirm all four components import cleanly by running:
python -c "from fl_topk.task import Net, load_data, train, test; print('All task functions imported')"

You should see All task functions imported with no errors.

✔️ Awesome, I've got everything!

Great. Double check you have saved task.py and that the import check passed.

ⓧ I'd like to double check the full code

Scroll up and confirm each of these sections is present in fl_topk/task.py, in this order:

  • Imports (torch, flwr_datasets, torchvision).
  • Net class with __init__ and forward methods.
  • fds = None followed by the load_data function.
  • train function.
  • test function.
Implement the Flower client

The client app defines what happens on each simulated IoT device during a federated round. Flower calls the client's training handler to send it the global model, let it train locally, and collect the updated weights.

  • Open fl_topk/client_app.py and add the imports and app declaration by pasting the following code:
import torch
from flwr.client import ClientApp
from flwr.common import Message, RecordDict, ArrayRecord, MetricRecord

from fl_topk.task import Net, load_data, train, test

app = ClientApp()
  • Below the app = ClientApp() line, add the training handler that each client runs during a federated round:
@app.train()
def train_fn(msg: Message, context):
    """Train the model on local data and return updated parameters."""
    # Read training config from pyproject.toml
    run_config = context.run_config
    local_epochs = int(run_config["local-epochs"])
    lr = float(run_config["lr"])

    partition_id = context.node_config["partition-id"]
    num_partitions = context.node_config["num-partitions"]
    device = torch.device("cpu")
    model = Net().to(device)
    # Load the global model weights sent by the server
    global_arrays = msg.content["arrays"]
    model.load_state_dict(global_arrays.to_torch_state_dict())
    # Train on this client's local data partition
    trainloader, _ = load_data(partition_id, num_partitions)
    train_loss = train(model, trainloader, local_epochs, lr, device)
    # Package updated model weights to send back
    arrays = ArrayRecord()
    for k, v in model.state_dict().items():
        arrays[k] = v.cpu().numpy()

    metrics = MetricRecord(
        {
            "num-examples": float(len(trainloader.dataset)),
            "train-loss": train_loss,
        }
    )
    content = RecordDict({"arrays": arrays, "metrics": metrics})
    return Message(content=content, reply_to=msg)

What does the training handler do?

When the server starts a training round, Flower sends a Message containing the global model weights to each client. The @app.train() handler unpacks those weights, trains on the client's local data partition, then packages the updated weights into a new Message to send back.

The context.run_config dictionary reads values from your pyproject.toml (like local-epochs and lr). context.node_config tells each client which data partition it owns.

  • Save client_app.py.
  • Verify the client app imports correctly by running:
python -c "from fl_topk.client_app import app; print('ClientApp created')"

You should see ClientApp created printed.

Seeing an ImportError?

  • Check that your import line reads from flwr.client import ClientApp exactly.
  • Confirm Flower is installed in your active virtual environment by running flwr --version.

Still stuck? Help me fix my ClientApp import error

  • Now add the evaluation handler below train_fn in client_app.py:
@app.evaluate()
def evaluate_fn(msg: Message, context):
    """Evaluate the global model on local test data."""
    partition_id = context.node_config["partition-id"]
    num_partitions = context.node_config["num-partitions"]

    device = torch.device("cpu")
    model = Net().to(device)
    # Load global model weights and evaluate on test data
    global_arrays = msg.content["arrays"]
    model.load_state_dict(global_arrays.to_torch_state_dict())

    _, testloader = load_data(partition_id, num_partitions)
    loss, accuracy = test(model, testloader, device)
    # Return evaluation metrics to the server
    metrics = MetricRecord(
        {
            "num-examples": float(len(testloader.dataset)),
            "eval-loss": loss,
            "eval-accuracy": accuracy,
        }
    )
    content = RecordDict({"metrics": metrics})
    return Message(content=content, reply_to=msg)

What does the evaluation handler do?

The @app.evaluate() handler runs after aggregation each round. The server sends the updated global model to clients, who evaluate it on a test set and return accuracy and loss metrics. Unlike training, evaluation does not change the model weights.

  • Save client_app.py.

✔️ Awesome, I've got everything!

Great. Confirm you saved client_app.py with both the train_fn and evaluate_fn handlers.

ⓧ I'd like to double check the full code

Scroll up and confirm each of these sections is present in fl_topk/client_app.py, in this order:

  • Imports (torch, flwr.client, flwr.common, fl_topk.task).
  • app = ClientApp() declaration.
  • @app.train() decorated train_fn function.
  • @app.evaluate() decorated evaluate_fn function.
Implement the server and run the federation

The server app orchestrates the entire federation. It initializes the global model, sends it to clients each round, collects their updates, and aggregates them using the FedAvg strategy.

  • Open fl_topk/server_app.py and add the imports and app declaration:
import json
from pathlib import Path
from typing import Optional

import torch
from flwr.common import ArrayRecord, MetricRecord
from flwr.server import ServerApp
from flwr.serverapp.strategy import FedAvg

from fl_topk.task import Net, test

app = ServerApp()
  • Below app = ServerApp(), add the main function with the centralized evaluation callback:
@app.main()
def main(grid, context):
    """Run the federated learning experiment."""
    run_config = context.run_config
    num_rounds = int(run_config["num-server-rounds"])
    sparsity_ratio = float(run_config.get("sparsity-ratio", 0.0))
    # Track results per round for later visualization
    results_log = {
        "sparsity_ratio": sparsity_ratio,
        "rounds": [],
    }

    def evaluate_fn(
        server_round: int, arrays: ArrayRecord
    ) -> tuple[Optional[float], Optional[MetricRecord]]:
        """Evaluate the global model on the server side."""
        model = Net()
        model.load_state_dict(arrays.to_torch_state_dict())
        from fl_topk.task import load_data
        _, testloader = load_data(0, 10)
        loss, accuracy = test(model, testloader, torch.device("cpu"))
        print(f"Round {server_round}: accuracy={accuracy:.4f}, loss={loss:.4f}, sparsity={sparsity_ratio}")
        results_log["rounds"].append(
            {"round": server_round, "accuracy": accuracy, "loss": loss}
        )
        return loss, MetricRecord({"centralized-accuracy": accuracy})

What does evaluate_fn do?

The evaluate_fn callback runs on the server after each round of aggregation. It reconstructs the global model from the aggregated weights, evaluates it on a test set, and prints the accuracy. This gives you a round-by-round view of how the federation is learning.

It also appends each round's results to results_log, which gets saved to a JSON file. You will use these JSON files in Step 5 to plot your comparison chart.

  • Save server_app.py.
  • Verify the server module imports correctly by running:
python -c "from fl_topk.server_app import app; print('ServerApp created')"

You should see ServerApp created printed.

Seeing an ImportError for FedAvg?

  • Make sure the import line reads from flwr.serverapp.strategy import FedAvg. Older Flower tutorials use a different import path that was changed in recent versions.
  • Confirm Flower is installed in your active virtual environment by running flwr --version.

Still stuck? Help me fix my FedAvg import error in server_app.py

  • Now add the rest of the main function below the evaluate_fn callback. Make sure this code is indented at the same level as the results_log variable above:
    # Create initial model and convert to ArrayRecord
    model = Net()
    initial_arrays = ArrayRecord()
    for k, v in model.state_dict().items():
        initial_arrays[k] = v.cpu().numpy()

    # Run the federated learning rounds
    strategy = FedAvg()

    result = strategy.start(
        grid=grid,
        initial_arrays=initial_arrays,
        num_rounds=num_rounds,
        evaluate_fn=evaluate_fn,
    )

    # Save results to JSON for later analysis
    results_dir = Path("results")
    results_dir.mkdir(exist_ok=True)
    filename = f"results_sparsity_{sparsity_ratio}.json"
    with open(results_dir / filename, "w") as f:
        json.dump(results_log, f, indent=2)

    print(f"Results saved to results/{filename}")

What does this code do?

This creates a fresh CNN, wraps its weights in an ArrayRecord, and passes them to the FedAvg strategy as the starting global model. The strategy.start() call kicks off the federation: it sends the model to clients, collects their trained updates, averages the weights, and repeats for num_rounds rounds.

After all rounds complete, results are saved to a JSON file named by the sparsity ratio (results_sparsity_0.0.json for the baseline). You will collect multiple JSON files at different compression levels in Step 4.

  • Save server_app.py.

✔️ Awesome, I've got everything!

Great. Confirm you saved server_app.py with the full main function including evaluate_fn, model initialization, strategy start, and results saving.

ⓧ I'd like to double check the full code

Scroll up and confirm each of these sections is present in fl_topk/server_app.py, in this order:

  • Imports (json, Path, Optional, torch, flwr.common, flwr.server, flwr.serverapp.strategy, fl_topk.task).
  • app = ServerApp() declaration.
  • @app.main() decorated main function containing config reading, results_log, nested evaluate_fn, model init, strategy.start(), and results saving.

What about pyproject.toml?

Your pyproject.toml is already configured from Step 1. It tells Flower where to find your ServerApp and ClientApp, and sets the default run config: 5 server rounds, 2 local epochs, learning rate 0.01, and 10 simulated SuperNodes.

All three files are ready. Before you run the federation, think about what you expect to happen.

  • You have 10 clients with non-IID data, each training for 2 epochs per round, over 5 rounds total. Do you think accuracy will increase, decrease, or stay flat? What range do you expect for CIFAR-10 with a small CNN and only 5 rounds?
  • Launch the federated simulation by running:
flwr run .

The first run downloads the CIFAR-10 dataset, which may take a minute. After that, you will see output like this for each round:

Round 1: accuracy=0.2534, loss=2.0145, sparsity=0.0
Round 2: accuracy=0.3412, loss=1.8231, sparsity=0.0
Round 3: accuracy=0.3987, loss=1.7012, sparsity=0.0
Round 4: accuracy=0.4356, loss=1.6234, sparsity=0.0
Round 5: accuracy=0.4821, loss=1.5456, sparsity=0.0

Your exact numbers will differ, but accuracy should climb from around 20-30% in round 1 to roughly 40-55% by round 5. The federation is learning.

Why is accuracy only 40-55%?

CIFAR-10 has 10 classes, so random guessing gives 10%. Reaching 40-55% in just 5 rounds with a small CNN on non-IID data is solid progress. A centralized model trained for 50+ epochs on all the data could reach 75-85%, but federated training with heterogeneous data splits converges more slowly.

Seeing an error when running flwr run?

  • Make sure you are running the command from the directory that contains pyproject.toml (the parent of the fl_topk/ folder).
  • Confirm your virtual environment is active. Flower must be importable for the command to work.
  • If you see an import error mentioning fl_topk, check that fl_topk/__init__.py exists and that the [tool.flwr.app.components] paths in pyproject.toml match your file structure.

Still stuck? Help me debug my flwr run error

Your baseline federation is running and learning. Next up, you will add communication metering to find out exactly how many bytes each client transmits per round.

Measure Communication Cost Per Round

Your federation is training a CNN across 10 clients and printing accuracy each round. But you have no idea how much data each client actually sends back to the server.

For IoT devices with limited battery and bandwidth, communication cost is the bottleneck. Before you can optimize it, you need to measure it. In this step, you will add byte-counting utilities and wire them into the training pipeline so every round reports exactly how many bytes each client transmits.

In this step, get ready to:
  • Add byte-counting utility functions to task.py.
  • Update client_app.py to compute and report communication bytes after each training round.
  • Run the federation and observe the per-round communication cost.
Add byte-counting utilities to task.py

To measure communication cost, you need two functions: one that calculates the total bytes in a model update, and one that counts only the non-zero bytes. The second function will become critical once you start sparsifying updates in the next step.

  • Open fl_topk/task.py in your editor.
  • Scroll to the bottom of the file, after the test() function.
  • Add these two byte-counting functions at the end of the file:
def compute_bytes(state_dict):
    """Compute total bytes for a model state dict (float32)."""
    total = 0
    for tensor in state_dict.values():
        total += tensor.numel() * 4
    return total


def compute_effective_bytes(state_dict):
    """Compute effective bytes (non-zero elements only)."""
    total = 0
    for tensor in state_dict.values():
        total += torch.count_nonzero(tensor).item() * 4
    return total

What do these functions do?

compute_bytes() counts every parameter in the model state dict and multiplies by 4 (since each float32 value is 4 bytes). This gives the total size of the model update.

compute_effective_bytes() only counts non-zero parameters. Right now, every parameter has a value, so both functions return the same number. Once you apply sparsification in the next step, the difference becomes dramatic.

  • Save task.py.

✔️ Awesome, I've got everything!

Great. Make sure you have saved task.py before moving on.

ⓧ I'd like to double check the full code

Your task.py is too long to show in full here. Scroll through your file and confirm it contains these sections in order:

  • Imports at the top (torch, torch.nn, flwr_datasets, torchvision.transforms).
  • The Net class (CNN with 2 conv layers and 2 FC layers).
  • The load_data() function using DirichletPartitioner.
  • The train() and test() functions.
  • The two new functions: compute_bytes() and compute_effective_bytes() at the very bottom.
Wire byte tracking into client_app.py

The byte-counting functions exist, but nobody calls them yet. You need to import them into the client and use them after each training round to measure how many bytes the model update contains.

  • Open fl_topk/client_app.py in your editor.
  • Find this import line near the top of the file:
from fl_topk.task import Net, load_data, train, test
  • Replace it with this updated import that includes the two new functions:
from fl_topk.task import Net, load_data, train, test, compute_bytes, compute_effective_bytes

Now you need to update the train_fn function to compute and report byte metrics after training.

  • Inside train_fn, find this section that starts right after the training call:
    train_loss = train(model, trainloader, local_epochs, lr, device)

    arrays = ArrayRecord()
    for k, v in model.state_dict().items():
        arrays[k] = v.cpu().numpy()

    metrics = MetricRecord(
        {
            "num-examples": float(len(trainloader.dataset)),
            "train-loss": train_loss,
        }
    )
  • Replace that entire section with the following code:
    train_loss = train(model, trainloader, local_epochs, lr, device)

    # Measure communication cost of the model update
    local_state_dict = model.state_dict()
    total_bytes = compute_bytes(local_state_dict)
    effective_bytes = compute_effective_bytes(local_state_dict)

    arrays = ArrayRecord()
    for k, v in local_state_dict.items():
        arrays[k] = v.cpu().numpy()

    metrics = MetricRecord(
        {
            "num-examples": float(len(trainloader.dataset)),
            "train-loss": train_loss,
            "total-bytes": float(total_bytes),
            "effective-bytes": float(effective_bytes),
        }
    )

What changed here?

Instead of calling model.state_dict() twice, you now capture it once in local_state_dict. This lets you measure its byte size before packaging it into the ArrayRecord.

The two new metrics, total-bytes and effective-bytes, are included in the MetricRecord that gets sent back to the server. Flower automatically aggregates these across all clients each round.

  • Also update the docstring on the train_fn function. Find this line:
    """Train the model on local data and return updated parameters."""
  • Replace it with:
    """Train the model on local data and report communication metrics."""
  • Save client_app.py.

Seeing an import error later?

Double check that the import line in client_app.py includes both compute_bytes and compute_effective_bytes. The names must match the functions you added to task.py exactly.

Also confirm that you saved both task.py and client_app.py. An unsaved file means the changes are invisible to the Flower runtime.

Help me fix an import error in my Flower client

✔️ Awesome, I've got everything!

Great. Make sure you have saved client_app.py before moving on.

ⓧ I'd like to double check the full code

Your client_app.py is too long to show in full here. Scroll through the file and confirm it contains these pieces in order:

  • Imports including compute_bytes and compute_effective_bytes from fl_topk.task.
  • The train_fn function that captures local_state_dict, computes total_bytes and effective_bytes, and includes both in the MetricRecord.
  • The evaluate_fn function, unchanged from Step 2.
Run the federation and observe communication cost

Both files are updated. Time to run the federation again and see what the byte metrics reveal about your system's communication overhead.

Before you run, think about this: the CNN model has roughly 62,500 float32 parameters. How many bytes do you think a single client transmits per round?

  • Run the federation by running this command:
flwr run .

You should see per-round accuracy output just like before, with the federation completing 5 rounds. The byte metrics from each client are now being collected by Flower's aggregation pipeline.

What do the numbers tell you?

Your model has approximately 62,506 parameters. Each one is a float32 value (4 bytes). That means every client sends roughly 250,024 bytes (~250 KB) per round.

With 10 clients, the total bytes transmitted per round is about 2.5 MB. Over 5 rounds, that is roughly 12.5 MB total. For a small CNN on a desktop, that is trivial. For a battery-powered IoT sensor on a cellular connection, that is a serious energy drain.

Notice that total-bytes and effective-bytes are identical right now. Every parameter is non-zero, so each client sends the entire model. That is the problem you will fix in the next step.

Federation not completing?

If you see a NameError mentioning compute_bytes or compute_effective_bytes, the import in client_app.py is missing those names. Check the import line matches the updated version above.

If you see an AttributeError about local_state_dict, make sure you added the local_state_dict = model.state_dict() line before the byte-counting calls and the ArrayRecord loop.

Help me debug my Flower federation run

You can now see the problem in hard numbers: every client transmits the full model every round. The next step introduces Top-K gradient sparsification to slash those bytes dramatically.

Implement Top-K Gradient Sparsification

In the previous step, you measured the communication cost of your federation and saw that each round transmits roughly 2.5 MB across all 10 clients. Every client sends the full model every round.

For a battery-powered IoT device, that is a serious drain. Top-K gradient sparsification fixes this by keeping only the largest update values and zeroing out the rest before transmission. In this step, you will build the sparsification function, wire it into the client, and watch the byte count plummet.

In this step, get ready to:
  • Add a top_k_sparsify function to task.py that keeps only the top K% of update values per layer.
  • Update client_app.py to compute model deltas and apply sparsification before sending updates.
  • Run experiments at multiple sparsity levels and observe the accuracy-vs-byte trade-off.
Add Top-K sparsification logic to task.py

The idea behind Top-K sparsification is simple. Instead of sending every single parameter update, only send the values with the largest magnitude. Zero out everything else.

What is Top-K sparsification?

Imagine you trained a model locally and got 62,000 updated parameter values. Most of those changes are tiny and contribute little to learning. Top-K keeps only the top (1 - sparsity_ratio) fraction by absolute magnitude and sets the rest to zero.

At 90% sparsity, you transmit only the 10% most impactful values. The effective byte count drops by roughly 90%.

  • In task.py, scroll to the bottom of the file (below the compute_effective_bytes function) and add the first part of the top_k_sparsify function:
def top_k_sparsify(delta_dict, sparsity_ratio):
    """Apply Top-K sparsification to a model update delta.

    Keeps only the top (1 - sparsity_ratio) fraction of values by magnitude
    per layer and zeros out the rest.

    Returns:
        sparse_delta: dict of sparsified tensors
        effective_bytes: bytes needed to transmit non-zero values
        total_bytes: bytes for the full uncompressed delta
    """
    sparse_delta = {}
    total_params = 0
    nonzero_params = 0

    for key, delta in delta_dict.items():
        num_elements = delta.numel()
        total_params += num_elements

        if sparsity_ratio <= 0.0 or num_elements <= 1:
            sparse_delta[key] = delta.clone()
            nonzero_params += num_elements
            continue

What does this code do?

  • The function takes a dictionary of model update deltas (the difference between local and global parameters) and a sparsity_ratio between 0.0 and 1.0.
  • It loops through each layer in the model. If sparsity is zero or the layer has only one element, it keeps everything unchanged.
  • The continue skips to the next layer when no sparsification is needed.
  • Directly below the continue line (still inside the for loop), add the core sparsification logic and return statement:
        k = max(1, int(num_elements * (1.0 - sparsity_ratio)))
        flat = delta.flatten()
        _, topk_indices = torch.topk(flat.abs(), k)
        mask = torch.zeros_like(flat)
        mask[topk_indices] = 1.0
        sparse_delta[key] = (flat * mask).reshape(delta.shape)
        nonzero_params += k

    effective_bytes = nonzero_params * 4
    total_bytes = total_params * 4
    return sparse_delta, effective_bytes, total_bytes

What does this code do?

  • k = max(1, int(num_elements * (1.0 - sparsity_ratio))) calculates how many values to keep. At 90% sparsity, this keeps 10% of values. The max(1, ...) ensures at least one value survives per layer.
  • torch.topk(flat.abs(), k) finds the indices of the K largest values by absolute magnitude. A binary mask is created, and multiplying the delta by the mask zeroes out everything except the top-K values.
  • The function returns both effective_bytes (non-zero values times 4 bytes each) and total_bytes (the full uncompressed size) so you can compare them.
  • Save task.py.
  • Verify the function loads without errors by running this command:
python -c "from fl_topk.task import top_k_sparsify; print('top_k_sparsify loaded successfully')"

You should see top_k_sparsify loaded successfully printed in the terminal.

Seeing an ImportError or SyntaxError?

  • Check that the indentation of the second code block lines up inside the for loop. The k = max(...) line should be at the same indent level as the if sparsity_ratio block above it.
  • Make sure the effective_bytes and return lines are outside the for loop (aligned with sparse_delta = {}).
  • Confirm you saved the file before running the import test.

Still stuck? Help me debug my top_k_sparsify function in task.py. I am getting an import or syntax error after adding the function.

Modify client_app.py to apply sparsification

Right now, the client sends its entire trained model directly. To apply sparsification, the client needs to compute the delta (the difference between the trained model and the global model), sparsify that delta, then reconstruct the parameters before sending.

  • In client_app.py, find the import line at the top of the file:
from fl_topk.task import Net, load_data, train, test, compute_bytes, compute_effective_bytes
  • Replace compute_effective_bytes with top_k_sparsify so the line now looks like this:
from fl_topk.task import Net, load_data, train, test, top_k_sparsify, compute_bytes
  • Next, replace the entire body of the train_fn function (everything between the @app.train() decorator and the @app.evaluate() decorator) with the following code. Start with this first section, which reads config, loads the model, saves a copy of the global parameters, trains, and computes the delta:
@app.train()
def train_fn(msg: Message, context):
    """Train the model on local data and return sparsified updates."""
    # Read config
    run_config = context.run_config
    local_epochs = int(run_config["local-epochs"])
    lr = float(run_config["lr"])
    sparsity_ratio = float(run_config.get("sparsity-ratio", 0.0))

    # Get partition info
    partition_id = context.node_config["partition-id"]
    num_partitions = context.node_config["num-partitions"]

    device = torch.device("cpu")
    model = Net().to(device)

    # Load global model from message
    global_arrays = msg.content["arrays"]
    global_state_dict = global_arrays.to_torch_state_dict()
    model.load_state_dict(global_state_dict)

    # Save a copy of global params for delta computation
    global_params = {k: v.clone() for k, v in global_state_dict.items()}

    # Load data and train
    trainloader, _ = load_data(partition_id, num_partitions)
    train_loss = train(model, trainloader, local_epochs, lr, device)

    # Compute delta (local - global)
    local_state_dict = model.state_dict()
    delta = {k: local_state_dict[k] - global_params[k] for k in local_state_dict}

What changed from the previous version?

  • sparsity_ratio is now read from the run config, defaulting to 0.0 (no compression).
  • The global model parameters are stored in global_params before training. After training, the code computes a delta: the element-wise difference between the trained model and the global model. This delta is what gets sparsified.
  • Directly below the delta line you just added, paste the rest of train_fn which applies sparsification and builds the reply:
    # Apply Top-K sparsification
    sparse_delta, effective_bytes, total_bytes = top_k_sparsify(delta, sparsity_ratio)

    # Reconstruct parameters: global + sparse_delta
    updated_params = {
        k: global_params[k] + sparse_delta[k] for k in global_params
    }

    # Build reply
    arrays = ArrayRecord()
    for k, v in updated_params.items():
        arrays[k] = v.cpu().numpy()

    metrics = MetricRecord(
        {
            "num-examples": float(len(trainloader.dataset)),
            "train-loss": train_loss,
            "effective-bytes": float(effective_bytes),
            "total-bytes": float(total_bytes),
        }
    )
    content = RecordDict({"arrays": arrays, "metrics": metrics})
    return Message(content=content, reply_to=msg)

What does this code do?

  • top_k_sparsify(delta, sparsity_ratio) applies the sparsification you wrote in task.py. It returns the sparse delta along with the byte counts.
  • The client reconstructs the full parameters by adding the sparse delta back to the global model (global_params[k] + sparse_delta[k]). This means the server receives a model that only reflects the top-K changes.
  • The effective-bytes metric now reflects the sparsified count (only non-zero values), so the server logs will show the real savings.
  • Make sure the evaluate_fn function below train_fn is unchanged from the previous step.
  • Save client_app.py.
  • Verify nothing is broken by running the federation with default sparsity (0.0):
flwr run .

You should see accuracy improving over 5 rounds, just like before. The byte counts should also be similar to your baseline because sparsity-ratio defaults to 0.0 (no compression).

Seeing an error when running the federation?

  • If you see an ImportError mentioning top_k_sparsify, confirm you saved task.py and that the function name matches exactly.
  • If you see an error about compute_effective_bytes not found, you still have the old import line. Make sure you replaced it with the new one that imports top_k_sparsify.
  • If accuracy is at 0% or 10% every round, check that the delta computation and reconstruction are correct. The reconstructed parameters should be global_params[k] + sparse_delta[k], not just sparse_delta[k] alone.

Help me debug my client_app.py after adding Top-K sparsification. The federation is failing or producing unexpected results.

Run experiments at different sparsity levels

Your federation now supports configurable sparsification. The --run-config flag lets you override the sparsity-ratio value from pyproject.toml without editing any files.

What does --run-config do?

Flower reads default config values from the [tool.flwr.app.config] section of pyproject.toml. The --run-config flag overrides those values for a single run. This makes it easy to sweep through different hyperparameters without changing code.

  • Run the 50% sparsity experiment by running this command:
flwr run . --run-config "sparsity-ratio=0.5"

Watch the terminal output. You should see the accuracy still improving, and the results are saved to results/results_sparsity_0.5.json.

Now for the big test. Before you run the next command, what do you think will happen to accuracy when you drop 90% of all update values?

  • Run the 90% sparsity experiment by running this command:
flwr run . --run-config "sparsity-ratio=0.9"

The effective bytes should drop by roughly 90% compared to the baseline. But notice the accuracy. It drops noticeably compared to the 0% and 50% runs.

Why does accuracy drop?

When you zero out 90% of the update values, those small contributions are permanently lost. Over multiple rounds, this information loss compounds. The model converges to a weaker solution because important gradient information is being discarded every round.

This accuracy-communication trade-off is the central tension in communication-efficient federated learning. The Secret Mission addresses this with a technique called error feedback.

  • Run the 95% sparsity experiment by running this command:
flwr run . --run-config "sparsity-ratio=0.95"

At 95% sparsity, accuracy degrades further. Only 5% of update values are transmitted each round.

  • Note the final accuracy and sparsity for each of your four runs (0%, 50%, 90%, 95%). You will use these numbers in the next step to build a comparison chart.
  • Confirm that the results files were saved by running:
ls results/

You should see JSON files for each sparsity level, including results_sparsity_0.0.json, results_sparsity_0.5.json, results_sparsity_0.9.json, and results_sparsity_0.95.json.

Missing a results file?

  • If results_sparsity_0.0.json is missing, re-run the baseline with no config override:
flwr run .
  • If any other file is missing, re-run the corresponding experiment. Each run saves its own file automatically.

Help me figure out why my Flower experiment results are not being saved to the results directory.

✔️ Awesome, I've got everything!

Great. Double check you have saved both task.py and client_app.py.

ⓧ I'd like to double check the full code

Both files are longer than 30 lines, so scroll up and verify the following sections are present and in order.

In task.py:

  • The Net class, load_data, train, test, compute_bytes, and compute_effective_bytes are all unchanged from the previous step.
  • At the bottom, the new top_k_sparsify function appears with its docstring, the edge-case check (if sparsity_ratio <= 0.0), the torch.topk mask logic, and the return sparse_delta, effective_bytes, total_bytes at the end.

In client_app.py:

  • The import line reads from fl_topk.task import Net, load_data, train, test, top_k_sparsify, compute_bytes (no compute_effective_bytes).
  • train_fn reads sparsity_ratio from config, saves global_params, computes a delta, calls top_k_sparsify, and reconstructs updated_params before building the reply.
  • evaluate_fn is unchanged from the previous step.

You now have hard numbers showing that Top-K sparsification slashes communication cost. But the accuracy trade-off is hard to evaluate from raw terminal output alone. Next up, you will build a chart that plots accuracy against total bytes transmitted across all four sparsity levels.

Compare Strategies and Visualize Trade-offs

In the previous step, you ran experiments at four different sparsity levels and watched accuracy and byte counts shift with each run. You now have four JSON result files sitting in the results/ directory, each capturing per-round accuracy and loss for a specific sparsity ratio.

Raw numbers scattered across separate terminal runs are hard to compare. A single visualization that plots accuracy against communication cost reveals the Pareto frontier. This is where you find the sweet spot between maximum byte savings and minimal accuracy loss, exactly the kind of analysis you would present in a paper.

In this step, get ready to:
  • Create a plotting script that loads experiment results from JSON files.
  • Generate a scatter chart comparing accuracy against communication bytes across sparsity levels.
Create plot_results.py to visualize the trade-off

The four result files your server saved each contain per-round accuracy and loss data for one sparsity level. You need a script that reads all of them, computes total effective bytes per experiment, and plots the trade-off on a single chart.

  • Create a new file called plot_results.py in your project root (the same folder that contains pyproject.toml) and paste the following code:
import json
from pathlib import Path

import matplotlib.pyplot as plt


def load_results(results_dir="results"):
    """Load all experiment results from JSON files."""
    results = []
    results_path = Path(results_dir)

    # Check if results directory exists
    if not results_path.exists():
        print(f"No results directory found at {results_dir}")
        print("Run experiments first with different sparsity-ratio values.")
        return results

    # Read each experiment's JSON file
    for filepath in sorted(results_path.glob("results_sparsity_*.json")):
        with open(filepath) as f:
            data = json.load(f)
            results.append(data)
    return results

What does this code do?

load_results() scans the results/ directory for files matching the pattern results_sparsity_*.json. It reads each one and returns a list of dictionaries, one per experiment.

sorted() ensures the files load in a consistent order (0.0, 0.5, 0.9, 0.95) regardless of filesystem ordering.

  • Save plot_results.py.
  • Verify the data loading works by running this command:
python -c "from plot_results import load_results; r = load_results(); print(f'Loaded {len(r)} experiment results')"

You should see Loaded 4 experiment results printed in your terminal.

Seeing 0 results or a file not found error?

  • Make sure you are running the command from the project root directory (the one containing pyproject.toml and the results/ folder).
  • Confirm your result files exist by running ls results/. You should see four JSON files from your Step 4 experiments.

Still stuck? Ask for help.

  • Now add the plot_accuracy_vs_communication() function below load_results() in plot_results.py. This first section extracts data points from each experiment:
def plot_accuracy_vs_communication(results):
    """Plot accuracy vs communication cost for different sparsity levels."""
    if not results:
        print("No results to plot.")
        return

    fig, ax = plt.subplots(1, 1, figsize=(10, 6))

    sparsity_levels = []
    accuracies = []
    comm_costs = []
    # Total bytes for one full model update (62,506 float32 params)
    model_bytes = 62506 * 4

    for result in results:
        sparsity = result["sparsity_ratio"]
        rounds_data = result["rounds"]
        if not rounds_data:
            continue

        final_accuracy = rounds_data[-1]["accuracy"]
        num_rounds = len(rounds_data)

        # Effective bytes = only the non-zero fraction, across 10 clients
        effective_ratio = 1.0 - sparsity
        total_effective_bytes = model_bytes * effective_ratio * 10 * num_rounds

        sparsity_levels.append(sparsity)
        accuracies.append(final_accuracy)
        comm_costs.append(total_effective_bytes / (1024 * 1024))

What does this code do?

For each experiment, the function reads the final round's accuracy and calculates total effective bytes transmitted across all 10 clients and all 5 rounds.

model_bytes = 62506 * 4 is the full model size in bytes (62,506 float32 parameters, each 4 bytes). The effective_ratio multiplier reflects how much of the model each client actually transmits. At 90% sparsity, only 10% of values are non-zero, so effective_ratio = 0.1.

  • Save plot_results.py.
  • Confirm the function is valid by running this import check:
python -c "from plot_results import plot_accuracy_vs_communication; print('Function loaded successfully')"

You should see Function loaded successfully. This confirms the data extraction code has no syntax errors.

  • Still inside the plot_accuracy_vs_communication() function, add the scatter plot and sparsity labels directly below the comm_costs.append line:
    # Plot each experiment as a data point
    ax.scatter(comm_costs, accuracies, s=120, zorder=5, color="steelblue")

    # Label each point with its sparsity ratio
    for i, sparsity in enumerate(sparsity_levels):
        label = f"  sparsity={sparsity:.0%}"
        if sparsity == 0.0:
            label = "  baseline (no compression)"
        ax.annotate(
            label,
            (comm_costs[i], accuracies[i]),
            fontsize=9,
            ha="left",
            va="bottom",
        )

What does this code do?

ax.scatter() draws one dot per experiment on the chart, with communication cost on the x-axis and accuracy on the y-axis.

The annotation loop adds a text label next to each dot. The baseline (0% sparsity) gets a special label so it stands out as the reference point.

Analyze the results and save the chart

The scatter plot shows where each sparsity level lands, but you also need to quantify how much communication each level saves relative to the baseline. Red annotations showing percentage byte reductions make the trade-off immediately readable.

  • Still inside the same function, add the byte-reduction annotations, axis configuration, and save logic below the label annotation loop. Then add the if __name__ block at the bottom of the file (outside the function):
    # Show percentage byte reduction relative to baseline
    if comm_costs and accuracies:
        baseline_bytes = max(comm_costs)
        for i, sparsity in enumerate(sparsity_levels):
            if sparsity > 0 and baseline_bytes > 0:
                reduction = (1 - comm_costs[i] / baseline_bytes) * 100
                ax.annotate(
                    f"  -{reduction:.0f}% bytes",
                    (comm_costs[i], accuracies[i]),
                    fontsize=8,
                    ha="left",
                    va="top",
                    color="red",
                )

    # Configure axis labels and title
    ax.set_xlabel("Total Effective Bytes Transmitted (MB)", fontsize=12)
    ax.set_ylabel("Final Test Accuracy", fontsize=12)
    ax.set_title(
        "Accuracy vs Communication Cost: Top-K Gradient Sparsification",
        fontsize=14,
    )
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig("accuracy_vs_communication.png", dpi=150)
    print("Chart saved to accuracy_vs_communication.png")
    plt.show()


if __name__ == "__main__":
    results = load_results()
    plot_accuracy_vs_communication(results)

What does this code do?

The byte-reduction block compares each sparsity level's communication cost against the baseline (0% sparsity) and annotates the percentage saving in red. This makes it easy to read exactly how much bandwidth each compression level saves.

plt.savefig() saves the chart as accuracy_vs_communication.png at 150 DPI, and plt.show() opens a window so you can inspect it interactively.

  • Save plot_results.py.

Before you run the script, think about what the chart will look like. You ran four experiments with sparsity ratios of 0.0, 0.5, 0.9, and 0.95. Do you expect the four data points to form a straight line, or do you think the accuracy drop will accelerate at higher sparsity levels?

  • Generate the chart by running:
python plot_results.py

A matplotlib window opens showing your scatter chart. You should see four labeled data points. The baseline sits in the top-right (high accuracy, high bytes), while the 95% sparsity point sits in the bottom-left (low accuracy, low bytes).

Close the chart window. Your terminal should print Chart saved to accuracy_vs_communication.png, and the PNG file now exists in your project root.

What does the chart tell you?

Look for the "knee" of the curve. This is the data point where you get the largest byte savings with the smallest accuracy penalty. Typically, 50% sparsity cuts communication cost in half with only a modest accuracy drop, while 90% and 95% sparsity deliver diminishing returns as accuracy degrades more steeply.

This is your Pareto frontier. Points along it represent the best possible accuracy for a given communication budget. Any strategy below the frontier wastes bytes without gaining accuracy.

Chart window not appearing?

  • If you are running over SSH or in a headless environment, plt.show() may not work. The script still saves the PNG file. Open accuracy_vs_communication.png from your file browser instead.
  • If you see "No results to plot", confirm your result files exist by running ls results/. You need the four JSON files from Step 4.

Need more help? Ask here.

✔️ Awesome, I've got everything!

Great. Double check you have saved plot_results.py and that accuracy_vs_communication.png exists in your project root.

ⓧ I'd like to double check the full code

plot_results.py is built from the following sections. Scroll up through your file and confirm each one is present, in order:

  • Imports (json, Path, matplotlib.pyplot).
  • load_results() function that reads JSON files from results/.
  • plot_accuracy_vs_communication() function containing data extraction, scatter plot, sparsity labels, byte-reduction annotations, axis labels, title, save, and show.
  • if __name__ == "__main__" block that calls load_results() then plot_accuracy_vs_communication().

You now have a quantitative chart that maps the accuracy-vs-communication frontier for Top-K gradient sparsification. Next up, the Secret Mission challenges you to recover the accuracy lost at high sparsity levels using error feedback.

Secret mission

Add Error Feedback to Recover Lost Accuracy

Your 90% sparsification cut communication by 90%, but accuracy took a hit. Error feedback recovers that lost accuracy by storing discarded update values and rolling them into the next round. This technique powers state-of-the-art FL compression methods like FedSparQ and TAMUNA.

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:

  • fl-topk-env virtual environment (contains Flower, PyTorch, and all pip-installed packages).
  • fl_topk/ directory with your Python source files (task.py, client_app.py, server_app.py).
  • Project configuration and scripts (pyproject.toml, plot_results.py).
  • results/ directory with experiment JSON files and accuracy_vs_communication.png chart.

Keep everything running

No action needed. Choose this if you are still actively experimenting with sparsity levels or extending the project.

  • Your virtual environment stays active and ready for more flwr run experiments.
  • All result JSON files and your chart remain in the results/ directory for further analysis.
  • The fl-topk-env folder uses roughly 1-2 GB of disk space (mostly PyTorch).

Pause - I'll come back to this later

Deactivate the virtual environment to free up your terminal. All project files stay on disk so you can reactivate and continue later.

  • Deactivate your virtual environment by running:
deactivate
  • When you are ready to resume, navigate back to your project folder and reactivate the environment.

macOS

source fl-topk-env/bin/activate

Windows

fl-topk-env\Scripts\activate

Delete - I don't want to use this again

Remove all project files, experiment results, and the virtual environment. This frees up disk space and leaves no trace of the project.

  • Deactivate your virtual environment first by running:
deactivate
  • Navigate to the parent of your project folder (your Desktop) and delete everything.

macOS

  • From inside your project folder, run these commands to delete all project files and the virtual environment:
cd ~/Desktop
rm -rf fl-topk-env fl_topk results plot_results.py pyproject.toml accuracy_vs_communication.png
  • Run ls ~/Desktop to confirm that fl_topk and fl-topk-env are no longer listed.

Windows

  • From inside your project folder, run these commands to delete all project files and the virtual environment:
cd ~\Desktop
rmdir /s /q fl-topk-env fl_topk results
del /q plot_results.py pyproject.toml accuracy_vs_communication.png
  • Run dir to confirm that fl_topk and fl-topk-env are no longer listed.

Nice Work!

Nice Work!

Nice work! You just built a complete federated learning experimental harness that quantifies the accuracy-vs-communication trade-off for gradient compression on simulated IoT clients.

You learned how to:

  • Simulate a 10-client FedAvg federation training a CNN on non-IID CIFAR-10 data using Flower and PyTorch, with per-round byte metering that reveals the communication bottleneck for IoT FL.
  • Implemented a Top-K gradient sparsification module that keeps only the largest update values by magnitude, reducing transmitted bytes by up to 95%.
  • Produced a quantitative accuracy-vs-communication trade-off chart across four sparsity levels, showing the Pareto frontier where byte savings start to cost too much accuracy.
  • Secret Mission: Added error feedback (residual accumulation) to recover accuracy lost from aggressive sparsification, a technique used in state-of-the-art FL compression research.

Ready to quiz yourself?