A walkthrough of a federated learning system hardened with differential privacy — clients train locally, only noisy model updates travel, and raw data never leaves the device.
Traditional ML pulls every client's data into one place to train a model. That single store becomes the most valuable — and most attractive — target on the network. Toggle between the two approaches below.
Each device trains on its own data and shares just the gradients — small, abstract numbers — with the server, which averages them into a global model. The raw records never leave home.
REDUCED BREACH RISKThere's a second, subtler problem federated learning alone doesn't solve: trained models memorize. Even without touching raw data directly, an attacker who studies a model's outputs or gradients can sometimes reconstruct facts about the training set. That's what differential privacy is for — jump to the privacy lab to see it in action.
Federated learning and differential privacy each block a different piece of the attack surface. Select an attack to see how the two techniques respond.
Unauthorized access to a centralized database exposing sensitive information — leading to financial loss, legal exposure, and reputational damage.
Federated learning keeps data decentralized on local devices instead of pooling it, removing the single high-value target. Differential privacy adds noise to any update that does leak, keeping it anonymized and low-value to an attacker.
The steps run in order — each one depends on the last — so they're numbered as a real sequence. Click through to see what each stage does and the code behind it.
The MNIST dataset (70,000 handwritten digits) is split into n subsets, one per simulated client — mimicking how real clients each hold their own local, possibly uneven, data. PyTorch DataLoaders manage batching for each partition.
from lib.federated import prepare_dataset trainloaders, valloaders, testloader = prepare_dataset( trainset=trainset, testset=testset, batch_size=32, num_partitions=3 )
start_server boots the coordinator using Flower (fl). Without a custom strategy it defaults to FedAvg — averaging every client's update, weighted by how many examples that client trained on.
from lib.federated import start_server def main(cfg): start_server( host_addr='localhost:5050', num_rounds=2 ) if __name__ == "__main__": main()
Any PyTorch model works. The test setup uses a small CNN — two convolutional layers with max-pooling, then three fully connected layers — with standard train/test functions built on cross-entropy loss.
class Net(nn.Module): def __init__(self, num_classes): self.conv1 = nn.Conv2d(1, 6, 5) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16*4*4, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, num_classes) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) # ...fully connected layers return x
fl_client wraps the model, optimizer, and that client's own data loaders into an object the Flower framework can drive — each client trains independently and only its resulting update leaves the device.
model = Net(num_classes=10) optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) client = fl_client( model=model, epochs=10, optimizer=optimizer, trainloader=trainloaders[int(client_id)], valloader=valloaders[int(client_id)], train=train, test=test, device=device ).to_client()
An Opacus PrivacyEngine wraps the model, optimizer, and data loader. noise_multiplier sets how much noise masks each gradient; max_grad_norm clips any single gradient so no one data point can dominate an update.
privacy_engine = PrivacyEngine()
model, optimizer, data_loader = privacy_engine.make_private(
module=model,
optimizer=optimizer,
data_loader=trainloaders[int(client_id)],
noise_multiplier=1.1,
max_grad_norm=1.0,
)start_client connects the prepared client to the server and kicks off training. A small shell script starts the server, then spins up all clients in parallel for a full federated round.
from lib.federated import fl_client, start_client start_client(server_address='localhost:5050', client=client) # run.sh — starts server + 3 clients python server.py & for i in $(seq 0 2); do python client.py --client-id $i & done wait
The pipeline above is what the code does. This is how the codebase — ronnie-1947/federated-learning — is actually structured to make that possible, and the two or three decisions that make it work.
federated_tutorial/ ├── run.sh boots 1 server + 3 clients ├── server.py Hydra entrypoint → starts Flower server ├── client.py entrypoint, --client-id N ├── model.py CNN + plain train/test loops ├── dataset.py MNIST download helper ├── cuda.py GPU availability check └── lib/ ├── federated.py Flower client/server glue ├── data.py prepare_dataset() partitioner └── diff-privacy.py Opacus wrapping helper
The codebase is split so that neither half knows the other exists:
Swap in a different PyTorch model and dataset, and the federated/DP machinery keeps working untouched — that's the whole point of building it as a library rather than a one-off script.
flwr_clientFlower can't call an arbitrary PyTorch model directly — it needs every client to speak one interface: NumPyClient. lib/federated.py implements exactly four methods to bridge that gap. Click through them.
Flower's server only understands lists of NumPy arrays — not PyTorch tensors or a state_dict. This method walks the model's state_dict() and detaches every tensor to .cpu().numpy(), so the weights can travel over the wire to the server for aggregation.
def get_parameters(self, config): # model weights → plain NumPy, ready to serialize return [val.cpu().numpy() for _, val in self.model.state_dict().items()]
The reverse trip: when the server broadcasts the newly averaged global weights, this rebuilds a PyTorch OrderedDict by zipping the model's own parameter names back onto the incoming NumPy arrays, then loads it with strict=True — a mismatch here fails loudly instead of silently.
def set_parameters(self, parameters): self.model.train() 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)
Called once per round. It loads the latest global weights, trains locally for the configured number of epochs using whatever train() callback was handed in — the CNN's own function, with zero Flower-specific code inside it — then hands the updated weights back along with how many examples this client trained on, which is what makes weighted FedAvg possible.
def fit(self, parameters, config): self.set_parameters(parameters=parameters) train_func(model=self.model, trainloader=self.trainloader, optimizer=self.optim, epochs=self.epochs, device=self.device, cb=self.train) return self.get_parameters(config={}), len(self.trainloader), {}
Also called once per round, on the client's held-out validation split. It loads the current global weights and runs the same test() callback used outside federation, returning loss, sample count, and accuracy — which the server's weighted_average function combines across every client, proportional to how much data each one holds.
def evaluate(self, parameters, config): self.set_parameters(parameters) loss, accuracy = test_func(model=self.model, testloader=self.valloader, device=self.device, cb=self.test) return float(loss), len(self.valloader), {"accuracy": accuracy}
These are the exact defaults in the repo — the fastest way to feel the privacy/accuracy trade-off is re-running with noise_multiplier=0.4 vs 2.0 and comparing what each client reports.
| Parameter | Default | Where | Controls |
|---|---|---|---|
num_partitions | 3 | client.py | How many simulated clients the data is split into |
num_rounds | 4 | server.py | Rounds of train → aggregate → broadcast |
epochs | 15 | client.py | Local epochs each client runs per round |
noise_multiplier | 1.1 | client.py | Higher = more privacy noise, lower utility |
max_grad_norm | 1.0 | client.py | Per-sample gradient clipping threshold |
lr / momentum | 0.1 / 0.9 | client.py | Local SGD optimizer settings |
# install pip install torch torchvision flwr opacus hydra-core numpy # server + 3 clients in one go cd federated_tutorial chmod +x run.sh ./run.sh
# terminal 1 python server.py # terminals 2, 3, 4 python client.py --client-id 0 python client.py --client-id 1 python client.py --client-id 2
Watched as separate processes, each client only ever prints its own local training/eval logs — never anyone else's data — which is the easiest way to actually see what "federated" means at runtime.
Swap the dataset by replacing dataset.py and the Net class — nothing in lib/ assumes MNIST. Swap the aggregation strategy by passing a custom strategy into start_server() instead of the default FedAvg. Simulate non-IID clients by adjusting the partition ratio in prepare_dataset() so clients hold unequal shares of the data.
This is the core idea behind Opacus's noise_multiplier. Each dot is one client's gradient update. Drag the slider — more noise scrambles individual points (harder to trace back to any one person) but also blurs the signal the model learns from.
Opacus adds Gaussian noise scaled by this value to every clipped gradient before it's aggregated — mathematically bounding how much any single data point can influence the model.
Trained with a batch size of 32, learning rate 0.1, momentum 0.9, and 100 local epochs per client per round. Hover any point for exact values.
| num_round | 4 |
| num_clients | 10 |
| min_available_clients | 10 |
| batch_size | 32 |
| local_epochs | 100 |
Loss fell steadily from 43.68 to 39.62 across four rounds, and accuracy climbed from 4.77% to 11.65%. The trend is healthy but the numbers are modest — the report is explicit that more rounds and tuning would be needed for a production-grade model.
Federated learning multiplies compute cost across every simulated client — the available hardware wasn't enough to train at real scale.
Sparse docs for libraries like PySyft meant testing several frameworks before settling on the Flower + Opacus combination used here.