Tutorial

Federated Learning for Medical Imaging using Flower and PyTorch

Difficulty: Advanced Time: 25 min read

Introduction

In healthcare, data silos are a massive bottleneck. Hospitals cannot easily share patient MRI or CT scans due to strict privacy regulations (HIPAA/GDPR). Federated Learning (FL) solves this by training algorithms across multiple decentralized edge devices or servers holding local data samples, without exchanging them. In this tutorial, we will use the open-source framework Flower (flwr) with PyTorch to train a federated medical image classification model.

Architecture Overview


graph LR
    Server(Central Flower Server) --> Client1(Hospital A Client)
    Server --> Client2(Hospital B Client)
    Server --> Client3(Hospital C Client)
    
    Client1 -.-> |Local Training on MRI Data| Client1
    Client2 -.-> |Local Training on MRI Data| Client2
    Client3 -.-> |Local Training on MRI Data| Client3
    
    Client1 --> |Upload Model Weights| Server
    Client2 --> |Upload Model Weights| Server
    Client3 --> |Upload Model Weights| Server
    
    Server -.-> |FedAvg Aggregation| Server
    Server --> |Broadcast Updated Global Model| Client1
    Server --> |Broadcast Updated Global Model| Client2
    Server --> |Broadcast Updated Global Model| Client3

Prerequisites

  • Python 3.9+
  • Flower (flwr)
  • PyTorch, Torchvision, Opacus (for Differential Privacy)
  • Basic understanding of client-server architectures

Step 1: Install Dependencies

pip install flwr torch torchvision opacus

Step 2: Defining the PyTorch Model and Training Logic

First, we define a standard PyTorch CNN (like a simplified ResNet or DenseNet) and the local training and evaluation functions. This code will be run by all clients locally.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import FakeData # Using FakeData for simulation, replace with your MRI dataset
from torchvision.transforms import ToTensor

# Simple CNN for demonstration
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 53 * 53, 120) # Assuming 224x224 input
        self.fc2 = nn.Linear(120, 2) # Binary classification (e.g., Tumor / No Tumor)

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

def train(net, trainloader, epochs):
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    for _ in range(epochs):
        for images, labels in trainloader:
            optimizer.zero_grad()
            loss = criterion(net(images), labels)
            loss.backward()
            optimizer.step()

def test(net, testloader):
    criterion = nn.CrossEntropyLoss()
    correct, total, loss = 0, 0, 0.0
    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.dataset), correct / total

Step 3: Creating the Flower Client

We wrap the PyTorch code inside a Flower NumPyClient. This client tells the framework how to get parameters from the model, train the model locally, and evaluate it.

import flwr as fl
from collections import OrderedDict

class FlowerClient(fl.client.NumPyClient):
    def __init__(self, net, trainloader, testloader):
        self.net = net
        self.trainloader = trainloader
        self.testloader = testloader

    def get_parameters(self, config):
        return [val.cpu().numpy() for _, val in self.net.state_dict().items()]

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

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

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

Step 4: Running the Server with Advanced Aggregation (FedProx)

While standard Federated Averaging (FedAvg) works well, healthcare data is often highly non-IID (Independent and Identically Distributed). One hospital might scan mostly older patients, while another scans mostly pediatric patients. To account for this data heterogeneity, we can use the FedProx aggregation strategy on the server.

# Start the server using FedProx to handle statistical heterogeneity across hospitals
python -c "
import flwr as fl
strategy = fl.server.strategy.FedProx(
    fraction_fit=1.0, 
    fraction_evaluate=1.0, 
    min_fit_clients=2, 
    min_evaluate_clients=2, 
    min_available_clients=2,
    proximal_mu=0.1 # The proximal term hyperparameter
)
fl.server.start_server(
    server_address='0.0.0.0:8080', 
    config=fl.server.ServerConfig(num_rounds=5),
    strategy=strategy
)"

Then, in separate terminals (representing different hospitals), start the clients:

# Start Client 1 (Hospital A)
python -c "import flwr as fl; from your_module import FlowerClient, SimpleCNN; client = FlowerClient(SimpleCNN(), train_loader_1, test_loader_1); fl.client.start_numpy_client(server_address='127.0.0.1:8080', client=client)"

# Start Client 2 (Hospital B)
python -c "import flwr as fl; from your_module import FlowerClient, SimpleCNN; client = FlowerClient(SimpleCNN(), train_loader_2, test_loader_2); fl.client.start_numpy_client(server_address='127.0.0.1:8080', client=client)"

Conclusion

Federated Learning with Flower allows diverse healthcare institutions to collaborate on building robust, generalizable AI models while keeping sensitive patient data strictly on-premises. By integrating advanced strategies like FedProx and Differential Privacy libraries (like Opacus), we can ensure models are both highly accurate across diverse populations and strictly mathematically private.