wellallyTech

Posted on Sep 12

Privacy-First Healthcare: Building a Skin Rash Classifier with Federated Learning and Flower 🌸

#ai #python #opensource #tutorial

In the world of medical AI, we face a massive paradox: we need massive datasets to build accurate models, but medical data is—rightfully—locked behind strict privacy walls like HIPAA and GDPR. How do we train a skin rash classification model across multiple hospitals without ever seeing a single patient's photo?

Enter Federated Learning (FL). 🛡️

Today, we are diving deep into the Flower (flwr) framework and PyTorch to demonstrate how "learning in public" (while keeping data private) is changing the game for Healthcare AI and Data Privacy. By the end of this guide, you’ll understand how to orchestrate a decentralized training session where the model travels, but the data stays home.

The Architecture: How Federated Learning Works

In traditional Machine Learning, we pull all data into a central server. In Federated Learning, we do the opposite. We send the model to the data, train it locally, and only send the "knowledge" (weights/gradients) back to the server.

graph TD
    subgraph "Central Server (Aggregator)"
        A[Global Model v1] --> B{Strategy: FedAvg}
        B -->|Sends Weights| C[Client 1: Hospital A]
        B -->|Sends Weights| D[Client 2: Hospital B]
        B -->|Sends Weights| E[Client 3: Clinic C]
    end

    subgraph "Edge Devices (Local Training)"
        C -->|Local Gradients| B
        D -->|Local Gradients| B
        E -->|Local Gradients| B
    end

    subgraph "Data Privacy Layer"
        F[Rash Images - Private] -.-> C
        G[Rash Images - Private] -.-> D
        H[Rash Images - Private] -.-> E
    end

Prerequisites 🛠️

To follow along, you'll need a basic understanding of PyTorch and the following stack:

Python 3.9+
Flower (flwr): The lightweight federated learning framework.
PyTorch: For building our CNN classifier.

pip install flwr torch torchvision

Step 1: Defining the Skin Rash Classifier

We'll start with a simple Convolutional Neural Network (CNN) tailored for image classification. In a real-world scenario, you might use a ResNet-50, but let's keep it lightweight for this demo.

import torch
import torch.nn as nn
import torch.nn.functional as F

class RashNet(nn.Module):
    def __init__(self):
        super(RashNet, 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 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 3) # 3 Classes: Eczema, Psoriasis, Healthy

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

Step 2: The Flower Client Logic

The "Magic" happens in the FlowerClient. This class tells the server how to interact with the local data at the hospital's edge node.

import flwr as fl
from collections import OrderedDict

class SkinCareClient(fl.client.NumPyClient):
    def __init__(self, model, trainloader, valloader):
        self.model = model
        self.trainloader = trainloader
        self.valloader = valloader

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

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

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        # Standard PyTorch Training Loop (simplified)
        train(self.model, self.trainloader, epochs=1)
        return self.get_parameters(config={}), len(self.trainloader.dataset), {}

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

The "Official" Way: Leveling Up Your AI Strategy 🥑

While this demo covers the basics of FL, deploying these systems in production requires handling heterogeneous data distributions (Non-IID), differential privacy, and secure multi-party computation.

For a deep dive into production-grade AI architectures and advanced privacy-preserving patterns, I highly recommend checking out the technical deep-dives at WellAlly Blog. It’s a fantastic resource for developers looking to move beyond "Hello World" into scalable, secure AI systems.

Step 3: Starting the Federated Server

Finally, we need a "conductor" to manage the training rounds and aggregate the weights using an algorithm like FedAvg.

import flwr as fl

# Define the strategy (Federated Averaging)
strategy = fl.server.strategy.FedAvg(
    fraction_fit=1.0,  # Sample 100% of available clients for training
    min_fit_clients=2, # Wait for at least 2 clients to connect
    min_available_clients=2,
)

# Start the Flower server
fl.server.start_server(
    server_address="0.0.0.0:8080",
    config=fl.server.ServerConfig(num_rounds=3),
    strategy=strategy,
)

Why This Matters 🚀

Data Sovereignty: The hospitals never lose control of their data.
Reduced Latency: No need to upload multi-gigabyte datasets to the cloud.
Collaborative Intelligence: Small clinics with limited data can benefit from the global model trained across large research hospitals.

Conclusion

Federated Learning is not just a buzzword; it’s a necessity for the future of ethical AI. Using Flower and PyTorch, we can build systems that respect user privacy while still pushing the boundaries of medical science.

Are you ready to decentralize your models? Drop a comment below if you’ve worked with FL before or if you’re planning to implement it in your next project!

Don't forget to visit wellally.tech/blog for more advanced tutorials on AI privacy and edge computing! 💻✨

DEV Community