Building and training the model

Now the model. We build a small 2D U-Net from scratch — deliberately, so that every value it learns lives in named_parameters() and federates with the three-line extract/reload from Concepts.

Why GroupNorm, not BatchNorm

One design choice matters for federated learning: we normalise with GroupNorm, not BatchNorm.

BatchNorm keeps running statistics (mean and variance) as buffers — and buffers are not in named_parameters(). The simple "extract the parameters, average them, reload" approach would silently leave those statistics behind, so each site's BatchNorm would drift out of sync with the averaged weights. GroupNorm has no running buffers — its scale and shift are ordinary parameters — so the entire model state travels when we federate. (This is the practical side of what a weight is: parameters federate, buffers don't.)

The U-Net

import torch
import torch.nn as nn


def gn(c):
    return nn.GroupNorm(num_groups=min(8, c), num_channels=c)


class DoubleConv(nn.Module):
    def __init__(self, cin, cout):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(cin, cout, 3, padding=1), gn(cout), nn.ReLU(inplace=True),
            nn.Conv2d(cout, cout, 3, padding=1), gn(cout), nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.net(x)


class UNet2D(nn.Module):
    def __init__(self, in_ch=1, out_ch=1, base=16):
        super().__init__()
        self.pool = nn.MaxPool2d(2)
        self.d1 = DoubleConv(in_ch, base)
        self.d2 = DoubleConv(base, base * 2)
        self.d3 = DoubleConv(base * 2, base * 4)
        self.bott = DoubleConv(base * 4, base * 8)
        self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
        self.u3 = DoubleConv(base * 8, base * 4)
        self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
        self.u2 = DoubleConv(base * 4, base * 2)
        self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
        self.u1 = DoubleConv(base * 2, base)
        self.out = nn.Conv2d(base, out_ch, 1)

    def forward(self, x):
        c1 = self.d1(x)
        c2 = self.d2(self.pool(c1))
        c3 = self.d3(self.pool(c2))
        b = self.bott(self.pool(c3))
        x = self.u3(torch.cat([self.up3(b), c3], dim=1))
        x = self.u2(torch.cat([self.up2(x), c2], dim=1))
        x = self.u1(torch.cat([self.up1(x), c1], dim=1))
        return self.out(x)  # raw logits

The loss

Lesions are tiny, so we combine a Dice term (good for imbalanced segmentation) with binary cross-entropy:

def dice_loss(logits, target, eps=1.0):
    prob = torch.sigmoid(logits)
    num = 2 * (prob * target).sum(dim=(1, 2, 3)) + eps
    den = prob.sum(dim=(1, 2, 3)) + target.sum(dim=(1, 2, 3)) + eps
    return 1 - (num / den).mean()


def dice_score(logits, target, thr=0.5, eps=1.0):
    pred = (torch.sigmoid(logits) > thr).float()
    num = 2 * (pred * target).sum(dim=(1, 2, 3)) + eps
    den = pred.sum(dim=(1, 2, 3)) + target.sum(dim=(1, 2, 3)) + eps
    return (num / den).mean().item()

Train locally first

Before federating anything, confirm the model learns on one site's data. On CPU this is slow — keep the subset small, or use a GPU.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = UNet2D().to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
bce = nn.BCEWithLogitsLoss()


def run_epoch(loader, train=True):
    model.train(train)
    tot_loss, tot_dice, n = 0.0, 0.0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        with torch.set_grad_enabled(train):
            logits = model(x)
            loss = bce(logits, y) + dice_loss(logits, y)
            if train:
                opt.zero_grad(); loss.backward(); opt.step()
        bs = x.size(0)
        tot_loss += loss.item() * bs
        tot_dice += dice_score(logits, y) * bs
        n += bs
    return tot_loss / max(n, 1), tot_dice / max(n, 1)


for ep in range(3):
    tl, td = run_epoch(train_loader, train=True)
    vl, vd = run_epoch(val_loader, train=False)
    print(f"epoch {ep + 1}: train loss {tl:.3f} dice {td:.3f} | val loss {vl:.3f} dice {vd:.3f}")

You should see Dice climb over the epochs — even on a small subset. It won't be a great model yet (that needs more data and epochs), but it proves the pipeline works.

This is a stepping stone, not a benchmark

A 16-subject, 3-epoch run on CPU produces a weak model on purpose — the point is to get the federated loop working on something you can train in minutes. Scale up the subject count and epochs once the federation is running.

Next

Once the model trains sensibly on one site, wire it into BranchKey so several sites train together:

➡️ Federating it with BranchKey