Tutorial

Federated Learning for Medical Imaging with Flower and PyTorch

Difficulty: Advanced Time: 25 min read

Introduction

In healthcare, patient data is highly sensitive and often legally prohibited from leaving the hospital's internal network (e.g., due to HIPAA or GDPR). Federated Learning (FL) solves this by bringing the model to the data, rather than the data to the model. In FL, multiple hospitals (clients) train a model locally on their private data, and only the updated model weights are sent to a central server to be aggregated. In this tutorial, we will use the open-source FL framework Flower (flwr) and PyTorch to train a CNN across distributed clients.

Architecture Overview


graph TD
    Server((Central Aggregation Server))
    
    Server -- 1. Distribute Global Model --> Client1[Hospital A Client]
    Server -- 1. Distribute Global Model --> Client2[Hospital B Client]
    Server -- 1. Distribute Global Model --> Client3[Hospital C Client]
    
    Client1 -. 2. Train on Local MRI Data .-> Client1
    Client2 -. 2. Train on Local MRI Data .-> Client2
    Client3 -. 2. Train on Local MRI Data .-> Client3
    
    Client1 -- 3. Send Weight Updates --> Server
    Client2 -- 3. Send Weight Updates --> Server
    Client3 -- 3. Send Weight Updates --> Server
    
    Server -. 4. Aggregate Weights: FedAvg .-> Server

Prerequisites

  • Python 3.8+
  • PyTorch, torchvision
  • Flower (flwr)
  • Basic understanding of client-server architecture

Step 1: Install Dependencies

pip install flwr torch torchvision numpy

Step 2: Defining the Central Server

The server is responsible for coordinating the clients, sending them the global model, and aggregating their updates using an algorithm like FedAvg (Federated Averaging). Create a file named server.py.

import flwr as fl

# Define strategy: Federated Averaging
strategy = fl.server.strategy.FedAvg(
    fraction_fit=1.0,  # Fraction of clients used during training (1.0 = all connected clients)
    fraction_evaluate=0.5, # Fraction of clients used for evaluation
    min_fit_clients=2, # Minimum number of clients to be sampled for training
    min_evaluate_clients=2,
    min_available_clients=2, # Wait until at least 2 clients are connected before starting
)

# Start Flower server
print("Starting Central Aggregation Server...")
fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=fl.server.ServerConfig(num_rounds=3), # Run for 3 rounds of federated training
    strategy=strategy,
)

Step 3: Defining the Client Node with Non-IID Data

A major challenge in medical FL is that data is often Non-IID (Independent and Identically Distributed). For example, Hospital A might only see elderly patients, while Hospital B sees pediatric patients. We simulate this by heavily skewing the dataset partitioning. Create a file named client.py.

import flwr as fl
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import numpy as np
import sys

# 1. Define the PyTorch Model (e.g., a simple CNN for MNIST/MedMNIST)
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, 3, 1)
        self.conv2 = nn.Conv2d(16, 32, 3, 1)
        self.fc1 = nn.Linear(32 * 5 * 5, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 2. Define standard PyTorch train/test functions
def train(net, trainloader, epochs):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.01)
    net.train()
    for _ in range(epochs):
        for images, labels in trainloader:
            optimizer.zero_grad()
            criterion(net(images), labels).backward()
            optimizer.step()

def test(net, testloader):
    criterion = nn.CrossEntropyLoss()
    correct, total, loss = 0, 0, 0.0
    net.eval()
    with torch.no_grad():
        for images, labels in testloader:
            outputs = net(images)
            loss += criterion(outputs, labels).item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return loss / len(testloader), correct / total

# 3. Load LOCAL Private Data (Simulating Non-IID distribution)
# Run `python client.py 0` for client A, and `python client.py 1` for client B
client_id = int(sys.argv[1]) if len(sys.argv) > 1 else 0

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
full_trainset = datasets.MNIST("./data", train=True, download=True, transform=transform)
full_testset = datasets.MNIST("./data", train=False, download=True, transform=transform)

# Simulate Non-IID: Client 0 gets digits 0-4, Client 1 gets digits 5-9
if client_id == 0:
    idx = (full_trainset.targets <= 4)
else:
    idx = (full_trainset.targets > 4)

trainset = Subset(full_trainset, np.where(idx)[0])
trainloader = DataLoader(trainset, batch_size=32, shuffle=True)
testloader = DataLoader(full_testset, batch_size=32)

# Instantiate the model
net = SimpleCNN()

# 4. Define the Flower Client
class FlowerClient(fl.client.NumPyClient):
    def get_parameters(self, config):
        return [val.cpu().numpy() for _, val in net.state_dict().items()]

    def set_parameters(self, parameters):
        params_dict = zip(net.state_dict().keys(), parameters)
        state_dict = {k: torch.tensor(v) for k, v in params_dict}
        net.load_state_dict(state_dict, strict=True)

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        train(net, trainloader, epochs=1) # Train locally for 1 epoch
        return self.get_parameters(config={}), len(trainloader.dataset), {}

    def evaluate(self, parameters, config):
        self.set_parameters(parameters)
        loss, accuracy = test(net, testloader)
        return float(loss), len(testloader.dataset), {"accuracy": accuracy}

# Start Client
print(f"Connecting to server as Client {client_id}...")
fl.client.start_numpy_client(server_address="127.0.0.1:8080", client=FlowerClient())

Step 4: Running the Federation

To simulate the federated learning process, you will need to open multiple terminal windows.

# Terminal 1: Start the server
python server.py

# Terminal 2: Start Client A (Hospital 1 - Digits 0-4)
python client.py 0

# Terminal 3: Start Client B (Hospital 2 - Digits 5-9)
python client.py 1

As soon as both clients connect, the server will begin the FedAvg process. Even though Hospital 1 has never seen a digit > 4, and Hospital 2 has never seen a digit <= 4, the globally aggregated model will be able to recognize all 10 digits!

Conclusion

Federated Learning with frameworks like Flower enables multi-institutional medical research without compromising patient privacy. By aggregating model weights rather than raw data, hospitals can collaboratively build highly accurate models that are robust to diverse demographic distributions and non-IID data.