Skip to main content

Federating it with BranchKey

Now the point of the exercise: instead of training alone, each site trains locally and shares only its weights, which BranchKey averages and returns. Every site runs this same notebook against its own data and its own leaf — no raw data is ever exchanged.

Set up the entities​

In the dashboard (or via the API), create a tree, a branch, and a leaf for your site — see Getting Started. Save the leaf credentials JSON; the session_token is shown only once.

Configure the branch:

  • Aggregation service: average (standard federated averaging).
  • files_per_aggregation: the number of leaves that will upload each round.
Testing on your own first? Set files_per_aggregation = 1

An aggregation only fires once files_per_aggregation files have arrived. If you're running a single leaf but the branch expects 2, nothing happens — each upload waits forever. For a solo test set it to 1 (each upload aggregates immediately — the "average" of one is itself); raise it to your real leaf count when others join. This is the most common first-run mistake — see the FAQ.

Then start the run on the branch (status → start), or uploads will report "run is blocked".

Connect the client​

Save your leaf credentials to secret/leaf-1.json (the secret/ folder is gitignored), then:

import json, os
import numpy as np
from branchkey import Client, Credentials, APIConfig

CRED_PATH = os.environ.get("BK_CREDENTIALS", "secret/leaf-1.json")
with open(CRED_PATH) as f:
    credentials = Credentials.from_dict(json.load(f))

client = Client(
    credentials=credentials,
    api_config=APIConfig(host="https://app.branchkey.com"),
    use_websocket=True,   # recommended transport — port 443 only
)
print("connected — run status:", client.run_status)

One federated round​

A round is the loop from Concepts, applied to our model:

def federated_round(client, timeout=600):
    # 1. Train locally for one epoch.
    run_epoch(train_loader, train=True)

    # 2. Extract weights and upload. `weighting` = our number of training slices.
    weighting, params = client.convert_pytorch_numpy(
        model.named_parameters(), weighting=len(train_ds)
    )
    weights_file = client.save_weights("round_weights", weighting, params)
    file_id = client.file_upload(weights_file)
    print("uploaded:", file_id)

    # 3. Wait for the aggregated result and download it.
    agg_id = client.queue.get(block=True, timeout=timeout)
    agg_path = client.file_download(agg_id)

    # 4. Load the averaged weights back, in parameter order.
    npz = np.load(agg_path, allow_pickle=False)
    keys = sorted(npz.files, key=lambda k: int(k.split("_")[1]))
    for (name, p), k in zip(model.named_parameters(), keys):
        p.data = torch.from_numpy(npz[k]).to(device)
    npz.close()

    # 5. Report a metric for this round.
    vl, vd = run_epoch(val_loader, train=False)
    client.send_performance_metrics(agg_id, json.dumps({"dice": vd, "loss": vl}), "test")
    return vd


for r in range(3):
    dice = federated_round(client)
    print(f"round {r + 1} complete — val dice {dice:.3f}")

client.close()
Sort layers numerically on reload

Step 4 sorts the downloaded layers by their integer suffix — int(k.split("_")[1]) — not alphabetically. With plain string sorting, layer_10 would come before layer_2, so a model with ten or more layers (like this U-Net) would load its weights into the wrong parameters. Always sort numerically. (See Shapes, layers & the NPZ format.)

A successful run looks like:

uploaded: 0956ab69-…
downloaded file cc13917d-…
round 1 complete — val dice 0.358

The Dice should stay in line with your local training (not collapse), which confirms the averaged weights landed in the right parameters.

Going multi-site​

To see genuine federated averaging, run this same notebook on a second machine with a second leaf on the same branch, and set files_per_aggregation = 2. Each round now averages both sites' weights — every participant ends up with a model shaped by data it never saw.

Later: graduating to nnU-Net​

The official UCSF-BMSR benchmarks use nnU-Net v1. It delivers strong accuracy, but it is a full framework you drive from the command line rather than a model you own — so federating it means subclassing its trainer to expose and replace the network's weights between rounds, on top of its 3D pipeline and data-format requirements. Get comfortable with the federated loop on the hand-written 2D model first; then the jump to nnU-Net is about where the weights live, not what federation has to do with them.