/federated-pytorch-test

Federated learning with PyTorch (federated averaging and consensus optimization): with 'reduced' bandwidth

Primary LanguagePythonApache License 2.0Apache-2.0

federated-pytorch-test

We train CNN models without having access to the full dataset. The CIFAR10 dataset is used in all examples. The CNN models can be chosen from simpler models similar to PyTorch or Tensorflow demos and in another example, they are ResNet18. In all cases, we use only 1/3 of the data for training each CNN model. We also compare the performance of federated averaging and consensus optimization in training the three models, without sharing the training data between models. Note that we only pass a subset of parameters between the models, unlike in normal federated averaging or consensus. This reduces the bandwidth required enormously!

The stochastic LBFGS optimizer is provided with the code. Further details are given in this paper. Also see this introduction.

Files included are:

lbfgsnew.py: New LBFGS optimizer

simple_models.py: Relatively simple CNN models for CIFAR10, derived from PyTorch/Tensorflow demos

no_consensus_trio.py: Train 3 models using 1/3 of the training data for each model

federated_trio.py: Train 3 models using 1/3 of the data, but with federated averaging

consensus_admm_trio.py: Train 3 models using 1/3 of the data, but with consensus optimization

federated_trio_resnet.py: Train 3 ResNet18 models using 1/3 of the data, but with federated averaging

consensus_admm_trio_resnet.py: Train 3 ResNet18 models using 1/3 of the data, but with consensus optimization