Federated Averaging (FedAvg) in PyTorch

An unofficial implementation of FederatedAveraging (or FedAvg) algorithm proposed in the paper Communication-Efficient Learning of Deep Networks from Decentralized Data in PyTorch. (implemented in Python 3.9.2.)

Implementation points

• Exactly implement the models ('2NN' and 'CNN' mentioned in the paper) to have the same number of parameters written in the paper.
  • 2NN: TwoNN class in models.py; 199,210 parameters
  • CNN: CNN class in models.py; 1,663,370 parameters
• Exactly implement the non-IID data split.
  • Each client has at least two digits in case of using MNIST dataset.
• Implement multiprocessing of client update and client evaluation.
• Support TensorBoard for log tracking.

Requirements

• See requirements.txt

Configurations

• See config.yaml

Run

• python3 main.py

Results

MNIST

• Number of clients: 100 (K = 100)
• Fraction of sampled clients: 0.1 (C = 0.1)
• Number of rounds: 500 (R = 500)
• Number of local epochs: 10 (E = 10)
• Batch size: 10 (B = 10)
• Optimizer: torch.optim.SGD
• Criterion: torch.nn.CrossEntropyLoss
• Learning rate: 0.01
• Momentum: 0.9
• Initialization: Xavier

Table 1. Final accuracy and the best accuracy

Table 2. Final loss and the least loss

Figure 1. MNIST 2NN model accuracy (IID: top / non-IID: bottom)

Figure 2. MNIST CNN model accuracy (IID: top / non-IID: bottom)

TODO

• Do CIFAR experiment (CIFAR10 dataset) & large-scale LSTM experiment (Shakespeare dataset)
• Learning rate scheduling
• More experiments with other hyperparameter settings (e.g., different combinations of B, E, K, and C)