PFLlib: Personalized Federated Learning Algorithm Library

👏 We will change the license to Apache-2.0 in the next release.

Figure 1: An Example for FedAvg. You can create a scenario using generate_DATA.py and run an algorithm using main.py, clientNAME.py, and serverNAME.py. For a new algorithm, you only need to add new features in clientNAME.py and serverNAME.py.

🎯We create a user-friendly algorithm library and evaluation platform for those new to federated learning. Join us in expanding the FL community by contributing your algorithms, datasets, and metrics to this project.

🎯If you find our repository useful, please cite the following paper:

@article{zhang2023pfllib,
  title={PFLlib: Personalized Federated Learning Algorithm Library},
  author={Zhang, Jianqing and Liu, Yang and Hua, Yang and Wang, Hao and Song, Tao and Xue, Zhengui and Ma, Ruhui and Cao, Jian},
  journal={arXiv preprint arXiv:2312.04992},
  year={2023}
}

• 37 traditional FL (tFL) and personalized FL (pFL) algorithms, 3 scenarios, and 20 datasets.

• Some experimental results are avalible in its paper and here.

• Refer to this guide to learn how to use it.

• This library can simulate scenarios using the 4-layer CNN on Cifar100 for 500 clients on one NVIDIA GeForce RTX 3090 GPU card with only 5.08GB GPU memory cost.

• PFLlib primarily focuses on data (statistical) heterogeneity. For algorithms and an evaluation platform that address both data and model heterogeneity, please refer to our extended project Heterogeneous Federated Learning (HtFLlib).

• As we strive to meet diverse user demands, frequent updates to the project may alter default settings and scenario creation codes, affecting experimental results.

• Closed issues may help you a lot.

• When submitting pull requests, please provide sufficient instructions and examples in the comment box.

The origin of the statistical heterogeneity phenomenon is the personalization of users, who generate non-IID (not Independent and Identically Distributed) and unbalanced data. With statistical heterogeneity existing in the FL scenario, a myriad of approaches have been proposed to crack this hard nut. In contrast, the personalized FL (pFL) may take advantage of the statistically heterogeneous data to learn the personalized model for each user.

Thanks to @Stonesjtu, this library can also record the GPU memory usage for the model. Following FedCG, we also introduce the DLG (Deep Leakage from Gradients) attack and PSNR (Peak Signal-to-Noise Ratio) metric to evaluate the privacy-preserving ability of tFL/pFL algorithms (please refer to ./system/flcore/servers/serveravg.py for example). Now we can train on some clients and evaluate performance on other new clients by setting args.num_new_clients in ./system/main.py. Note that not all the tFL/pFL algorithms support this feature.

Algorithms with code (updating)

Traditional FL (tFL)

• FedAvg — Communication-Efficient Learning of Deep Networks from Decentralized Data AISTATS 2017

  Update-correction-based tFL

• SCAFFOLD - SCAFFOLD: Stochastic Controlled Averaging for Federated Learning ICML 2020

  Regularization-based tFL

• FedProx — Federated Optimization in Heterogeneous Networks MLsys 2020

• FedDyn — Federated Learning Based on Dynamic Regularization ICLR 2021

  Model-splitting-based tFL

• MOON — Model-Contrastive Federated Learning CVPR 2021

• FedLC — Federated Learning With Label Distribution Skew via Logits Calibration ICML 2022

  Knowledge-distillation-based tFL

• FedGen — Data-Free Knowledge Distillation for Heterogeneous Federated Learning ICML 2021

• FedNTD — Preservation of the Global Knowledge by Not-True Distillation in Federated Learning NeurIPS 2022

Personalized FL (pFL)

• FedMTL (not MOCHA) — Federated multi-task learning NeurIPS 2017

• FedBN — FedBN: Federated Learning on non-IID Features via Local Batch Normalization ICLR 2021

  Meta-learning-based pFL

• Per-FedAvg — Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach NeurIPS 2020

  Regularization-based pFL

• pFedMe — Personalized Federated Learning with Moreau Envelopes NeurIPS 2020

• Ditto — Ditto: Fair and robust federated learning through personalization ICML 2021

  Personalized-aggregation-based pFL

• APFL — Adaptive Personalized Federated Learning 2020

• FedFomo — Personalized Federated Learning with First Order Model Optimization ICLR 2021

• FedAMP — Personalized Cross-Silo Federated Learning on non-IID Data AAAI 2021

• FedPHP — FedPHP: Federated Personalization with Inherited Private Models ECML PKDD 2021

• APPLE — Adapt to Adaptation: Learning Personalization for Cross-Silo Federated Learning IJCAI 2022

• FedALA — FedALA: Adaptive Local Aggregation for Personalized Federated Learning AAAI 2023

  Model-splitting-based pFL

• FedPer — Federated Learning with Personalization Layers 2019

• LG-FedAvg — Think Locally, Act Globally: Federated Learning with Local and Global Representations 2020

• FedRep — Exploiting Shared Representations for Personalized Federated Learning ICML 2021

• FedRoD — On Bridging Generic and Personalized Federated Learning for Image Classification ICLR 2022

• FedBABU — Fedbabu: Towards enhanced representation for federated image classification ICLR 2022

• FedGC — Federated Learning for Face Recognition with Gradient Correction AAAI 2022

• FedCP — FedCP: Separating Feature Information for Personalized Federated Learning via Conditional Policy KDD 2023

• GPFL — GPFL: Simultaneously Learning Generic and Personalized Feature Information for Personalized Federated Learning ICCV 2023

• FedGH — FedGH: Heterogeneous Federated Learning with Generalized Global Header ACM MM 2023

• FedDBE — Eliminating Domain Bias for Federated Learning in Representation Space NeurIPS 2023

• FedCAC — Bold but Cautious: Unlocking the Potential of Personalized Federated Learning through Cautiously Aggressive Collaboration ICCV 2023

• PFL-DA — Personalized Federated Learning via Domain Adaptation with an Application to Distributed 3D Printing Technometrics 2023

  Knowledge-distillation-based pFL (more in HtFLlib)

• FedDistill (FD) — Communication-Efficient On-Device Machine Learning: Federated Distillation and Augmentation under Non-IID Private Data 2018

• FML — Federated Mutual Learning 2020

• FedKD — Communication-efficient federated learning via knowledge distillation Nature Communications 2022

• FedProto — FedProto: Federated Prototype Learning across Heterogeneous Clients AAAI 2022

• FedPCL (w/o pre-trained models) — Federated learning from pre-trained models: A contrastive learning approach NeurIPS 2022

• FedPAC — Personalized Federated Learning with Feature Alignment and Classifier Collaboration ICLR 2023

Datasets and scenarios (updating)

For the label skew scenario, we introduce 14 famous datasets: MNIST, EMNIST, Fashion-MNIST, Cifar10, Cifar100, AG News, Sogou News, Tiny-ImageNet, Country211, Flowers102, GTSRB, Shakespeare, and Stanford Cars, they can be easy split into IID and non-IID version. Since some codes for generating datasets such as splitting are the same for all datasets, we move these codes into ./dataset/utils/dataset_utils.py. In the non-IID scenario, 2 situations exist. The first one is the pathological non-IID scenario, the second one is the practical non-IID scenario. In the pathological non-IID scenario, for example, the data on each client only contains the specific number of labels (maybe only 2 labels), though the data on all clients contains 10 labels such as the MNIST dataset. In the practical non-IID scenario, Dirichlet distribution is utilized (please refer to this paper for details). We can input balance for the iid scenario, where the data are uniformly distributed.

For the feature shift scenario, we use 3 datasets that are widely used in Domain Adaptation: Amazon Review (fetch raw data from this site), Digit5 (fetch raw data from this site), and DomainNet.

For the real-world scenario, we also introduce 3 naturally separated datasets: Omniglot (20 clients, 50 labels), HAR (Human Activity Recognition) (30 clients, 6 labels), PAMAP2 (9 clients, 12 labels). For the details of datasets and FL algorithms in IoT, please refer to FL-IoT.

If you need another data set, just write another code to download it and then use the utils.

Examples for MNIST

• MNIST

  cd ./dataset
  # python generate_MNIST.py iid - - # for iid and unbalanced scenario
  # python generate_MNIST.py iid balance - # for iid and balanced scenario
  # python generate_MNIST.py noniid - pat # for pathological noniid and unbalanced scenario
  python generate_MNIST.py noniid - dir # for practical noniid and unbalanced scenario
  # python generate_MNIST.py noniid - exdir # for Extended Dirichlet strategy 
  

The output of python generate_MNIST.py noniid - dir

Number of classes: 10
Client 0         Size of data: 2630      Labels:  [0 1 4 5 7 8 9]
                 Samples of labels:  [(0, 140), (1, 890), (4, 1), (5, 319), (7, 29), (8, 1067), (9, 184)]
--------------------------------------------------
Client 1         Size of data: 499       Labels:  [0 2 5 6 8 9]
                 Samples of labels:  [(0, 5), (2, 27), (5, 19), (6, 335), (8, 6), (9, 107)]
--------------------------------------------------
Client 2         Size of data: 1630      Labels:  [0 3 6 9]
                 Samples of labels:  [(0, 3), (3, 143), (6, 1461), (9, 23)]
--------------------------------------------------

Client 3         Size of data: 2541      Labels:  [0 4 7 8]
                 Samples of labels:  [(0, 155), (4, 1), (7, 2381), (8, 4)]
--------------------------------------------------
Client 4         Size of data: 1917      Labels:  [0 1 3 5 6 8 9]
                 Samples of labels:  [(0, 71), (1, 13), (3, 207), (5, 1129), (6, 6), (8, 40), (9, 451)]
--------------------------------------------------
Client 5         Size of data: 6189      Labels:  [1 3 4 8 9]
                 Samples of labels:  [(1, 38), (3, 1), (4, 39), (8, 25), (9, 6086)]
--------------------------------------------------
Client 6         Size of data: 1256      Labels:  [1 2 3 6 8 9]
                 Samples of labels:  [(1, 873), (2, 176), (3, 46), (6, 42), (8, 13), (9, 106)]
--------------------------------------------------
Client 7         Size of data: 1269      Labels:  [1 2 3 5 7 8]
                 Samples of labels:  [(1, 21), (2, 5), (3, 11), (5, 787), (7, 4), (8, 441)]
--------------------------------------------------
Client 8         Size of data: 3600      Labels:  [0 1]
                 Samples of labels:  [(0, 1), (1, 3599)]
--------------------------------------------------
Client 9         Size of data: 4006      Labels:  [0 1 2 4 6]
                 Samples of labels:  [(0, 633), (1, 1997), (2, 89), (4, 519), (6, 768)]
--------------------------------------------------
Client 10        Size of data: 3116      Labels:  [0 1 2 3 4 5]
                 Samples of labels:  [(0, 920), (1, 2), (2, 1450), (3, 513), (4, 134), (5, 97)]
--------------------------------------------------
Client 11        Size of data: 3772      Labels:  [2 3 5]
                 Samples of labels:  [(2, 159), (3, 3055), (5, 558)]
--------------------------------------------------
Client 12        Size of data: 3613      Labels:  [0 1 2 5]
                 Samples of labels:  [(0, 8), (1, 180), (2, 3277), (5, 148)]
--------------------------------------------------
Client 13        Size of data: 2134      Labels:  [1 2 4 5 7]
                 Samples of labels:  [(1, 237), (2, 343), (4, 6), (5, 453), (7, 1095)]
--------------------------------------------------
Client 14        Size of data: 5730      Labels:  [5 7]
                 Samples of labels:  [(5, 2719), (7, 3011)]
--------------------------------------------------
Client 15        Size of data: 5448      Labels:  [0 3 5 6 7 8]
                 Samples of labels:  [(0, 31), (3, 1785), (5, 16), (6, 4), (7, 756), (8, 2856)]
--------------------------------------------------
Client 16        Size of data: 3628      Labels:  [0]
                 Samples of labels:  [(0, 3628)]
--------------------------------------------------
Client 17        Size of data: 5653      Labels:  [1 2 3 4 5 7 8]
                 Samples of labels:  [(1, 26), (2, 1463), (3, 1379), (4, 335), (5, 60), (7, 17), (8, 2373)]
--------------------------------------------------
Client 18        Size of data: 5266      Labels:  [0 5 6]
                 Samples of labels:  [(0, 998), (5, 8), (6, 4260)]
--------------------------------------------------
Client 19        Size of data: 6103      Labels:  [0 1 2 3 4 9]
                 Samples of labels:  [(0, 310), (1, 1), (2, 1), (3, 1), (4, 5789), (9, 1)]
--------------------------------------------------
Total number of samples: 70000
The number of train samples: [1972, 374, 1222, 1905, 1437, 4641, 942, 951, 2700, 3004, 2337, 2829, 2709, 1600, 4297, 4086, 2721, 4239, 3949, 4577]
The number of test samples: [658, 125, 408, 636, 480, 1548, 314, 318, 900, 1002, 779, 943, 904, 534, 1433, 1362, 907, 1414, 1317, 1526]

Saving to disk.

Finish generating dataset.

Models

• for MNIST and Fashion-MNIST

  1. Mclr_Logistic(1*28*28)
  2. LeNet()
  3. DNN(1*28*28, 100) # non-convex

• for Cifar10, Cifar100 and Tiny-ImageNet

  1. Mclr_Logistic(3*32*32)
  2. FedAvgCNN()
  3. DNN(3*32*32, 100) # non-convex
  4. ResNet18, AlexNet, MobileNet, GoogleNet, etc.

• for AG_News and Sogou_News

  1. LSTM()
  2. fastText() in Bag of Tricks for Efficient Text Classification
  3. TextCNN() in Convolutional Neural Networks for Sentence Classification
  4. TransformerModel() in Attention is all you need

• for AmazonReview

  1. AmazonMLP() in Curriculum manager for source selection in multi-source domain adaptation

• for Omniglot

  1. FedAvgCNN()

• for HAR and PAMAP

  1. HARCNN() in Convolutional neural networks for human activity recognition using mobile sensors

Environments

Install CUDA v11.6.

Install conda latest and activate conda.

conda env create -f env_cuda_latest.yaml # You may need to downgrade the torch using pip to match the CUDA version

How to start simulating (examples for FedAvg)

• Download this project to an appropriate location using git.

  git clone https://github.com/TsingZ0/PFLlib.git
  

• Create proper environments (see Environments).

• Build evaluation scenarios (see Datasets and scenarios (updating)).

• Run evaluation:

  cd ./system
  python main.py -data MNIST -m CNN -algo FedAvg -gr 2000 -did 0 # using the MNIST dataset, the FedAvg algorithm, and the 4-layer CNN model
  

Note: It is preferable to tune algorithm-specific hyper-parameters before using any algorithm on a new machine.

Practical situations

If you need to simulate FL under practical situations, which includes client dropout, slow trainers, slow senders, and network TTL, you can set the following parameters to realize it.

• -cdr: The dropout rate for total clients. The selected clients will randomly drop at each training round.
• -tsr and -ssr: The rates for slow trainers and slow senders among all clients. Once a client is selected as a "slow trainer"/"slow sender", for example, it will always train/send slower than the original one.
• -tth: The threshold for network TTL (ms).

Easy to extend

It is easy to add new algorithms and datasets to this library.

• To add a new dataset into this library, all you need to do is write the download code and use the utils which is similar to ./dataset/generate_MNIST.py (you can also consider it as the template).

• To add a new algorithm, you can utilize the class Server and class Client, which are wrote in ./system/flcore/servers/serverbase.py and ./system/flcore/clients/clientbase.py, respectively.

• To add a new model, just add it into ./system/flcore/trainmodel/models.py.

• If you have a new optimizer while training, please add it into ./system/flcore/optimizers/fedoptimizer.py

• The evaluation platform is also convenient for users to build a new platform for specific applications, such as our FL-IoT and HtFLlib.

Experimental results

If you are interested in the experimental results (e.g., the accuracy) of the above algorithms, you can find some results in our accepted FL papers (i.e., FedALA, FedCP, GPFL, and DBE) listed as follows that also use this library. Please note that this developing project may not be able to reproduce the results on these papers, since some basic settings may change due to the requests of the community. For example, we previously set shuffle=False in clientbase.py

@inproceedings{zhang2023fedala,
  title={Fedala: Adaptive local aggregation for personalized federated learning},
  author={Zhang, Jianqing and Hua, Yang and Wang, Hao and Song, Tao and Xue, Zhengui and Ma, Ruhui and Guan, Haibing},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  volume={37},
  number={9},
  pages={11237--11244},
  year={2023}
}

@inproceedings{Zhang2023fedcp,
  author = {Zhang, Jianqing and Hua, Yang and Wang, Hao and Song, Tao and Xue, Zhengui and Ma, Ruhui and Guan, Haibing},
  title = {FedCP: Separating Feature Information for Personalized Federated Learning via Conditional Policy},
  year = {2023},
  booktitle = {Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining}
}

@inproceedings{zhang2023gpfl,
  title={GPFL: Simultaneously Learning Global and Personalized Feature Information for Personalized Federated Learning},
  author={Zhang, Jianqing and Hua, Yang and Wang, Hao and Song, Tao and Xue, Zhengui and Ma, Ruhui and Cao, Jian and Guan, Haibing},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  pages={5041--5051},
  year={2023}
}

@inproceedings{zhang2023eliminating,
  title={Eliminating Domain Bias for Federated Learning in Representation Space},
  author={Jianqing Zhang and Yang Hua and Jian Cao and Hao Wang and Tao Song and Zhengui XUE and Ruhui Ma and Haibing Guan},
  booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
  year={2023},
  url={https://openreview.net/forum?id=nO5i1XdUS0}
}