Extracting & reloading weights

This is the practical page. You know what weights are and how they're shaped and packaged. Now: how do you get them out of your model to upload, and put the aggregated ones back in?

Every federated round is the same round trip:

train locally → EXTRACT weights → upload → (BranchKey aggregates) → download → RELOAD weights → repeat

"Extract" and "reload" are the two steps that touch your model. They are only a few lines each.

The idea, in plain NumPy

Stripped of any framework, extracting means: produce a list of NumPy arrays, in a fixed order. Reloading means: take a list of NumPy arrays back and write them into the model, in the same order.

# EXTRACT — build the list you hand to save_weights:
parameters = [array_for_layer_0, array_for_layer_1, ...]

# RELOAD — read the averaged list back into your model:
for layer_index, averaged_array in enumerate(aggregated_layers):
    put averaged_array into layer number `layer_index` of your model

The only thing that changes between frameworks is how you ask the model for its arrays and how you write them back. Here are the three common cases.

PyTorch

PyTorch exposes weights through model.named_parameters(). The client library ships a helper, convert_pytorch_numpy, that turns those into the NumPy list save_weights wants — and also returns your chosen weighting:

# EXTRACT
weighting, parameters = client.convert_pytorch_numpy(
    model.named_parameters(),
    weighting=len(train_dataset),   # number of local training samples
)
weights_file = client.save_weights("model_weights", weighting, parameters)
file_id = client.file_upload(weights_file)

Under the hood the helper simply does, for each parameter, param.detach().cpu().numpy() — so if you prefer, the manual version is equally valid:

parameters = [p.detach().cpu().numpy() for _, p in model.named_parameters()]

Reloading — iterate the model's parameters in the same order and copy each averaged array back in:

import numpy as np, torch

aggregated_file = client.file_download(aggregation_id)
npz = np.load(aggregated_file, allow_pickle=False)

# Sort layer_0, layer_1, ... by their number (not alphabetically).
layer_keys = sorted(npz.files, key=lambda k: int(k.split("_")[1]))

for (name, param), key in zip(model.named_parameters(), layer_keys):
    param.data = torch.from_numpy(npz[key])
npz.close()

Order must be stable

named_parameters() yields parameters in a consistent order for a given model definition, so extract and reload line up automatically — as long as you don't change the model architecture between the two steps. If you reorder or rename layers, you reorder the list.

Keras / TensorFlow

Keras gives you the whole list directly with get_weights() and takes it back with set_weights():

# EXTRACT
parameters = model.get_weights()                # already a list of NumPy arrays
weighting = len(x_train)
weights_file = client.save_weights("model_weights", weighting, parameters)
file_id = client.file_upload(weights_file)

# RELOAD
import numpy as np
aggregated_file = client.file_download(aggregation_id)
npz = np.load(aggregated_file, allow_pickle=False)
layer_keys = sorted(npz.files, key=lambda k: int(k.split("_")[1]))
model.set_weights([npz[k] for k in layer_keys])
npz.close()

set_weights expects the arrays in the same order and shapes get_weights produced — which is exactly what you uploaded, so it lines up.

scikit-learn

Linear models keep their parameters in named attributes such as coef_ and intercept_. You choose which attributes constitute "the model" and list them in a fixed order:

# EXTRACT  (order you pick here is the order you must reload in)
parameters = [model.coef_, model.intercept_]
weighting = len(X_train)
weights_file = client.save_weights("model_weights", weighting, parameters)
file_id = client.file_upload(weights_file)

# RELOAD
import numpy as np
aggregated_file = client.file_download(aggregation_id)
npz = np.load(aggregated_file, allow_pickle=False)
layer_keys = sorted(npz.files, key=lambda k: int(k.split("_")[1]))
model.coef_      = npz[layer_keys[0]]
model.intercept_ = npz[layer_keys[1]]
npz.close()

Choosing the weighting

weighting is the one extra number you provide. It controls how strongly your update pulls the shared average. Three common choices:

weighting = len(train_dataset)            # by sample count (most common, recommended)
weighting = 1                             # equal say for every participant
weighting = len(train_dataset) * val_acc  # down-weight a poorly-performing local model

Only relative magnitudes matter: a participant with weighting=2000 counts twice as much as one with weighting=1000. See Configuring a branch for how the aggregator uses it.

Common pitfalls

• Not detaching (PyTorch). param.numpy() on a tensor that still requires gradients will error. Use param.detach().cpu().numpy() (or the convert_pytorch_numpy helper).
• Wrong order on reload. If your reload loop visits layers in a different order than extraction, you'll write the right shapes into the wrong places and the model will silently get worse. Keep the order identical.
• Alphabetical layer sorting. sorted(npz.files) puts layer_10 after layer_1. Always sort by the integer suffix (shown in every snippet above).
• Changing architecture mid-federation. Everyone in a branch must share the same architecture, so the layer count and shapes match. Decide the architecture up front.