vaseline555 / Federated-Learning-in-PyTorch

Handy PyTorch implementation of Federated Learning (for your painless research)

Federated Learning in PyTorch

Implementations of various Federated Learning (FL) algorithms in PyTorch, especially for research purposes.

Implementation Details

Datasets

• Supports all image classification datasets in torchvision.datasets.
• Supports all text classification datasets in torchtext.datasets.
• Supports all datasets in LEAF benchmark (NO need to prepare raw data manually)
• Supports additional image classification datasets (TinyImageNet, CINIC10).
• Supports additional text classification datasets (BeerReviews).
• Supports tabular datasets (Heart, Adult, Cover).
• Supports temporal dataset (GLEAM)
• NOTE: don't bother to search raw files of datasets; the dataset can automatically be downloaded to the designated path by just passing its name!

Statistical Heterogeneity Simulations

• IID (i.e., statistical homogeneity)
• Unbalanced (i.e., sample counts heterogeneity)
• Pathological Non-IID (McMahan et al., 2016)
• Dirichlet distribution-based Non-IID (Hsu et al., 2019)
• Pre-defined (for datasets having natural semantic separation, including LEAF benchmark (Caldas et al., 2018))

Models

• LogReg (logistic regression), StackedTransformer (TransformerEncoder-based classifier)
• TwoNN, TwoCNN, SimpleCNN (McMahan et al., 2016)
• FEMNISTCNN, Sent140LSTM (Caldas et al., 2018))
• LeNet (LeCun et al., 1998), MobileNet (Howard et al., 2019), SqueezeNet (Iandola et al., 2016), VGG (Simonyan et al., 2014), ResNet (He et al., 2015)
• MobileNeXt (Daquan et al., 2020), SqueezeNeXt (Gholami et al., 2016), MobileViT (Mehta et al., 2021)
• DistilBERT (Sanh et al., 2019), SqueezeBERT (Iandola et al., 2020), MobileBERT (Sun et al., 2020)
• M5 (Dai et al., 2016)

Algorithms

• FedAvg and FedSGD (McMahan et al., 2016) Communication-Efficient Learning of Deep Networks from Decentralized Data
• FedAvgM (Hsu et al., 2019) Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification
• FedProx (Li et al., 2018) Federated Optimization in Heterogeneous Networks
• FedOpt (FedAdam, FedYogi, FedAdaGrad) (Reddi et al., 2020) Adaptive Federated Optimization

Evaluation schemes

• local: evaluate FL algorithm using holdout sets of (some/all) clients NOT participating in the current round. (i.e., evaluation of personalized federated learning setting)
• global: evaluate FL algorithm using global holdout set located at the server. (ONLY available if the raw dataset supports pre-defined validation/test set).
• both: evaluate FL algorithm using both local and global schemes.

Metrics

• Top-1 Accuracy, Top-5 Accuracy, Precision, Recall, F1
• Area under ROC, Area under PRC, Youden's J
• Seq2Seq Accuracy
• MSE, RMSE, MAE, MAPE
• $R^2$, $D^2$

Requirements

• See requirements.txt. (I recommend building an independent environment for this project, using e.g., Docker or conda)
• When you install torchtext, please check the version compatibility with torch. (See official repository)
• Plus, please install torch-related packages using one command provided by the official guide (See official installation guide); e.g., conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 torchtext==0.13.0 cudatoolkit=11.6 -c pytorch -c conda-forge

Configurations

• See python3 main.py -h.

Example Commands

• See shell files prepared in commands directory.

TODO

• Support another model, especially lightweight ones for cross-device FL setting. (e.g., EdgeNeXt)
• Support another structured dataset including temporal and tabular data, along with datasets suitable for cross-silo FL setting. (e.g., MedMNIST)
• Add other popular FL algorithms including personalized FL algorithms (e.g., SuPerFed).
• Attach benchmark results of sample commands.

Contact

Should you have any feedback, please create a thread in issue tab. Thank you :)