Federated Optimizers in PyTorch

What is Federated Learning?

Federated Learning is a distributed machine learning approach which enables model training on a large corpus of decentralized data.

From Towards Federated Learning at Scale: System Design

Federated Averaging

Description

• Communication-Efficient Learning of Deep Networks from Decentralized Data

Usage

optimizer = optimizers.FederatedAvgServer(model.parameters())
optimizer.zero_grad()
loss_fn(model(input), target).backward()

# On every client then do
optimizer_client = optimizers.FederatedAvgClient(model.parameters(), lr=0.1)
optimizer_client.zero_grad()
loss_fn(model(input), target).backward()
optimizer_client.step()
nk_grad = (n_training_examples, model.parameters())

# Send nk_grad from clients (1 to l) to the server
list_nk_grad = [nk_grad1, ..., nk_gradl]
optimizers.step(list_nk_grad)
# Redistribute updated model.parameters() from server to clients