Need help with FedMA?
Click the “chat” button below for chat support from the developer who created it, or find similar developers for support.

About the developer

IBM
223 Stars 59 Forks MIT License 4 Commits 9 Opened issues

Description

Code for Federated Learning with Matched Averaging, ICLR 2020.

Services available

!
?

Need anything else?

Contributors list

Federated Learning with Matched Averaging

This is the code accompanying the ICLR 2020 paper "Federated Learning with Matched Averaging " Paper link: [https://openreview.net/forum?id=BkluqlSFDS]

Overview


FedMA algorithm is designed for federated learning of modern neural network architectures e.g. convolutional neural networks (CNNs) and LSTMs. FedMA constructs the shared global model in a layer-wise manner by matching and averaging hidden elements (i.e. channels for convolution layers; hidden states for LSTM; neurons for fully connected layers) with similar feature extraction signatures.

Depdendencies


Tested stable depdencises: * python 3.6.5 (Anaconda) * PyTorch 1.1.0 * torchvision 0.2.2 * CUDA 10.0.130 * cuDNN 7.5.1 * lapsolver 1.0.2

Data Preparation


Language Models:

For the language model experiments, we used the Shakespeare dataset provided by project Leaf. Following the instructions to prepare Shakespeare dataset, we choose to use non-i.i.d., full-size dataset, and split 80% of the data points into the training dataset. Moreover, we set minimum number of samples per user at 9K. Thus, the following command returns our data partitioning:

./preprocess.sh -s niid --sf 1.0 -k 0 -t sample -tf 0.8 -k 9

Image Classification:

We simulate a heterogeneous partition for which batch sizes and class proportions are unbalanced. We simulate a heterogeneous partition by sampling proportion of the data points in each class across participating clients from a Dirichlet distribution. Due to the small concentration parameter (0.5) of the Dirichlet distribution, some sampled batches may not have any examples of certain classes of data. Details about this partition can be found in the

partition_data
function in
./utils.py
.

Experients over Language Task:


The source code involving language task experiments i.e. LSTM over the Shakespeare dataset locates in the folder

FedMA/language_modeling
. And we summarize the functionality of each script below.

| Script | Functionality | | ----------------------------- | ---------------------------------------- | |

ensemble_accuracy_calculator.py
| Evaluating the performance of ensemble accross local models trained on paritipating clients. | |
language_main.py
| Conducting
FedAvg
and
FedProx
experiments, which are used as baseline methods. | |
language_oneshot_matching.py
| Evaluating the performance of one-shot match i.e. PFNM-style model fusion. | |
language_whole_training.py
| Centralized training over one device i.e. we combine the local datasets and coduct centralized training. This is the strongest possible baseline for any Federated Leaarning method. | |
lstm_fedma_with_comm.py
| Our proposed "FedMA with communication algorithm". |

Experients over Image Classification Task:


The main result related to the image classification task i.e. VGG-9 on CIFAR-10 can be reproduced via running

./run.sh
. The following arguments to the
./main.py
file control the important parameters of the experiment.

| Argument | Description | | ----------------------------- | ---------------------------------------- | |

model
| The CNN architecture that each client train locally. | |
dataset
| Dataset to use. We use CIFAR-10 to study FedMA. | |
lr
| Inital learning rate that will be use. | |
retrain_lr
| The learning rate for the local re-training process. Usually set to the same value as
lr
| |
batch-size
| Batch size for the optimizers e.g. SGD or Adam. | |
epochs
| Locally training epochs. | |
retrain_epochs
| Local re-training epochs. | |
n_nets
| Number of participating local clients. | |
partition
| Data partitioning strategy. Set to
hetero-dir
for the simulated heterogeneous CIFAR-10 dataset. | |
comm_type
| Federated learning methods. Set to
fedavg
,
fedprox
, or
fedma
. | |
comm_round
| Number of communication rounds to use in
fedavg
,
fedprox
, and
fedma
. | |
retrain
| Flag to retrain the model or load from checkpoint. | |
rematching
| Flag to re-conduct the matching process or load from checkpoint. |

Sample command

python main.py --model=moderate-cnn \
--dataset=cifar10 \
--lr=0.01 \
--retrain_lr=0.01 \
--batch-size=64 \
--epochs=20 \
--retrain_epochs=20 \
--n_nets=16 \
--partition=hetero-dir \
--comm_type=fedma \
--comm_round=50 \
--retrain=True \
--rematching=True

Interpretability of FedMA:


The results of interpretability we presented in the FedMA paper are summerized in a jupyter notebook i.e.

./jupyter_notebook/Interpretability_fedma.ipynb
.

Handling Data Bias Experiments:


The handeling data bias experiments we presented the FedMA paper are summerized in the script

./dist_skew_main.py
. To reproduce the experiment, one can simply run:
bash run_dist_skew.sh

Sample command

python dist_skew_main.py --model=moderate-cnn \
--dataset=cifar10 \
--lr=0.01 \
--retrain_lr=0.01 \
--batch-size=64 \
--epochs=10 \
--retrain_epochs=20 \
--n_nets=2 \
--partition=homo \
--comm_type=fedma \
--retrain=True \
--rematching=True

Citing FedMA:


@inproceedings{
Wang2020Federated,
title={Federated Learning with Matched Averaging},
author={Hongyi Wang and Mikhail Yurochkin and Yuekai Sun and Dimitris Papailiopoulos and Yasaman Khazaeni},
booktitle={International Conference on Learning Representations},
year={2020},
url={https://openreview.net/forum?id=BkluqlSFDS}
}

We use cookies. If you continue to browse the site, you agree to the use of cookies. For more information on our use of cookies please see our Privacy Policy.