README.md

Benchmarking Federated Learning Methods.

Realizing Your Brilliant Ideas.

Having Fun with Federated Learning.

FL-bench welcomes PR on everything that can make this project better.

FL-bench 的简单介绍

Methods 🧬

Traditional FL Methods

• FedAvg -- Communication-Efficient Learning of Deep Networks from Decentralized Data (AISTATS'17)
• FedAvgM -- Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification (ArXiv'19)
• FedProx -- Federated Optimization in Heterogeneous Networks (MLSys'20)
• SCAFFOLD -- SCAFFOLD: Stochastic Controlled Averaging for Federated Learning (ICML'20)
• MOON -- Model-Contrastive Federated Learning (CVPR'21)
• FedDyn -- Federated Learning Based on Dynamic Regularization (ICLR'21)
• FedLC -- Federated Learning with Label Distribution Skew via Logits Calibration (ICML'22)
• FedGen -- Data-Free Knowledge Distillation for Heterogeneous Federated Learning (ICML'21)
• CCVR -- No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data (NeurIPS'21)
• FedOpt -- Adaptive Federated Optimization (ICLR'21)
• Elastic Aggregation -- Elastic Aggregation for Federated Optimization (CVPR'23)
• FedFed -- FedFed: Feature Distillation against Data Heterogeneity in Federated Learning (NeurIPS'23)

Personalized FL Methods

• pFedSim (My Work⭐) -- pFedSim: Similarity-Aware Model Aggregation Towards Personalized Federated Learning (ArXiv'23)
• Local-Only -- Local training only (without communication).
• FedMD -- FedMD: Heterogenous Federated Learning via Model Distillation (NeurIPS'19)
• APFL -- Adaptive Personalized Federated Learning (ArXiv'20)
• LG-FedAvg -- Think Locally, Act Globally: Federated Learning with Local and Global Representations (ArXiv'20)
• FedBN -- FedBN: Federated Learning On Non-IID Features Via Local Batch Normalization (ICLR'21)
• FedPer -- Federated Learning with Personalization Layers (AISTATS'20)
• FedRep -- Exploiting Shared Representations for Personalized Federated Learning (ICML'21)
• Per-FedAvg -- Personalized Federated Learning with Theoretical Guarantees: A Model-Agnostic Meta-Learning Approach (NeurIPS'20)
• pFedMe -- Personalized Federated Learning with Moreau Envelopes (NeurIPS'20)
• Ditto -- Ditto: Fair and Robust Federated Learning Through Personalization (ICML'21)
• pFedHN -- Personalized Federated Learning using Hypernetworks (ICML'21)
• pFedLA -- Layer-Wised Model Aggregation for Personalized Federated Learning (CVPR'22)
• CFL -- Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints (ArXiv'19)
• FedFomo -- Personalized Federated Learning with First Order Model Optimization (ICLR'21)
• FedBabu -- FedBabu: Towards Enhanced Representation for Federated Image Classification (ICLR'22)
• FedAP -- Personalized Federated Learning with Adaptive Batchnorm for Healthcare (IEEE'22)
• kNN-Per -- Personalized Federated Learning through Local Memorization (ICML'22)
• MetaFed -- MetaFed: Federated Learning among Federations with Cyclic Knowledge Distillation for Personalized Healthcare (IJCAI'22)
• FedRoD -- On Bridging Generic and Personalized Federated Learning for Image Classification (ICLR'22)
• FedProto -- FedProto: Federated prototype learning across heterogeneous clients (AAAI'22)
• FedPAC -- Personalized Federated Learning with Feature Alignment and Classifier Collaboration (ICLR'23)
• PeFLL -- PeFLL: Personalized Federated Learning by Learning to Learn (ICLR'24)
• FLUTE -- Federated Representation Learning in the Under-Parameterized Regime (ICML'24)

FL Domain Generalization Methods

• FedSR -- FedSR: A Simple and Effective Domain Generalization Method for Federated Learning (NeurIPS'22)
• ADCOL -- Adversarial Collaborative Learning on Non-IID Features (ICML'23)
• FedIIR -- Out-of-Distribution Generalization of Federated Learning via Implicit Invariant Relationships (ICML'23)

Environment Preparation 🧩

PyPI 🐍

pip install -r .env/requirements.txt

Poetry 🎶

For those China mainland users

poetry install --no-root -C .env

For others

cd .env && sed -i "10,14d" pyproject.toml && poetry lock --no-update && poetry install --no-root

Docker 🐳

For those China mainland users

docker pull registry.cn-hangzhou.aliyuncs.com/karhoutam/fl-bench:master

For others

docker pull ghcr.io/karhoutam/fl-bench:master

An example of building container

docker run -it --name fl-bench -v path/to/FL-bench:/root/FL-bench --privileged --gpus all ghcr.io/karhoutam/fl-bench:master

Easy Run 🏃‍♂️

ALL classes of methods are inherited from FedAvgServer and FedAvgClient. If you wanna figure out the entire workflow and detail of variable settings, go check src/server/fedavg.py and src/client/fedavg.py.

Step 1. Generate FL Dataset

Partition the MNIST according to Dir(0.1) for 100 clients

python generate_data.py -d mnist -a 0.1 -cn 100

About methods of generating federated dastaset, go check data/README.md for full details.

Step 2. Run Experiment

python main.py [--config-path, --config-name] [method=<METHOD_NAME> args...]

• method: The algorithm's name, e.g., method=fedavg. ❗ Method name should be identical to the .py file name in src/server.
• --config-path: Relative path to the directory of the config file. Defaults to config.
• --config-name: Name of .yaml config file (w/o the .yaml extension). Defaults to defaults, which points to config/defaults.yaml.

Such as running FedAvg with all defaults.

python main.py method=fedavg

Defaults are set in both config/defaults.yaml and src/utils/constants.py.

How To Customize FL method Arguments 🤖

• By modifying config file.
• By explicitly setting in CLI, e.g., python main.py --config-name my_cfg.yaml method=fedprox fedprox.mu=0.01.
• By modifying the default value in config/defaults.yaml or get_hyperparams() in src/server/<method>.py

⚠ For the same FL method argument, the priority of argument setting is CLI > Config file > Default value.

For example, the default value of fedprox.mu is 1,

# src/server/fedprox.py
class FedProxServer(FedAvgServer):

    @staticmethod
    def get_hyperparams(args_list=None) -> Namespace:
        parser = ArgumentParser()
        parser.add_argument("--mu", type=float, default=1.0)
        return parser.parse_args(args_list)

and your .yaml config file has

# config/your_config.yaml
...
fedprox:
  mu: 0.01

python main.py method=fedprox                                  # fedprox.mu = 1
python main.py --config-name your_config method=fedprox        # fedprox.mu = 0.01

Monitor 📈

FL-bench supports visdom and tensorboard.

Activate

👀 NOTE: You needs to launch visdom / tensorboard server by yourself.

# your_config.yaml
common:
  ...
  visible: tensorboard # options: [null, visdom, tensorboard]

Launch visdom / tensorboard Server

visdom

1. Run python -m visdom.server on terminal.
2. Go check localhost:8097 on your browser.

tensorboard

1. Run tensorboard --logdir=<your_log_dir> on terminal.
2. Go check localhost:6006 on your browser.

Parallel Training via Ray 🚀

This feature can vastly improve your training efficiency. At the same time, this feature is user-friendly and easy to use!!!

Activate (What You ONLY Need To Do)

# your_config.yaml
mode: parallel
parallel:
  num_workers: 2 # any positive integer that larger than 1
  ...
...

Manually Create Ray Cluster (Optional)

A Ray cluster would be created implicitly everytime you run experiment in parallel mode. Or you can create it manually by the command shown below to avoid creating and destroying cluster every time you run experiment.

ray start --head [OPTIONS]

👀 NOTE: You need to keep num_cpus: null and num_gpus: null in your config file for connecting to a existing Ray cluster.

# your_config_file.yaml
# Connect to an existing Ray cluster in localhost.
mode: parallel
parallel:
  ...
  num_gpus: null
  num_cpus: null
...

Arguments 🔧

FL-bench highly recommend through config file to customize your FL method and experiment settings.

FL-bench offers a default config file config/defaults.yaml that contains all required arguments and corresponding comments.

All common arguments have their default value. Go check config/defaults.yaml or DEFAULTS in src/utils/constants.py for all argument defaults.

NOTE :If your custom config file does not contain all required arguments, FL-bench will fill those missing arguments with their defaults that loaded from DEFAULTS.

About the default values of specific FL method arguments, go check corresponding src/server/<method>.py for the full details.

However, FL-bench also supports CLI arguments for quick changings. Here are some examples:

# Using config/defaults.yaml but change the method to FedProx and set its mu to 0.1.
python main.py method=fedprox fedprox.mu=0.1

# Change learning rate to 0.1.
python main.py optimizer.lr=0.1

# Change batch size to 128.
python main.py common.batch_size=128

Models 🤖

This benchmark supports bunch of models that common and integrated in Torchvision (check here for all):

• ResNet family
• EfficientNet family
• DenseNet family
• MobileNet family
• LeNet5
• ...

🤗 You can define your own custom model by filling the CustomModel class in src/utils/models.py and use it by defining model: custom in your .yaml config file.

Datasets and Partition Strategies 🎨

Regular Image Datasets

• MNIST (1 x 28 x 28, 10 classes)

• CIFAR-10/100 (3 x 32 x 32, 10/100 classes)

• EMNIST (1 x 28 x 28, 62 classes)

• FashionMNIST (1 x 28 x 28, 10 classes)

• Syhthetic Dataset

• FEMNIST (1 x 28 x 28, 62 classes)

• CelebA (3 x 218 x 178, 2 classes)

• SVHN (3 x 32 x 32, 10 classes)

• USPS (1 x 16 x 16, 10 classes)

• Tiny-ImageNet-200 (3 x 64 x 64, 200 classes)

• CINIC-10 (3 x 32 x 32, 10 classes)

Domain Generalization Image Datasets

• DomainNet (3 x ? x ?, 345 classes)
  • Go check data/README.md for the full process guideline 🧾.

Medical Image Datasets

• COVID-19 (3 x 244 x 224, 4 classes)

• Organ-S/A/CMNIST (1 x 28 x 28, 11 classes)

Customization Tips 💡

Implementing FL Method

The package() at server-side class is used for assembling all parameters server need to send to clients. Similarly, package() at client-side class is for parameters clients need to send back to server. You should always has super().package() in your override implementation.

• Consider to inherit your method classes from FedAvgServer and FedAvgClient for maximum utilizing FL-bench's workflow.

• You can also inherit your method classes from advanced methods, e.g., FedBN, FedProx, etc. Which will inherit all functions, variables and hyperparamter settings. If you do that, you need to careful design your method in order to avoid potential hyperparamters and workflow conflicts.

class YourServer(FedBNServer):
  ...

class YourClient(FedBNClient):
  ...

• For customizing your server-side process, consider to override the package() and aggregate().

• For customizing your client-side training, consider to override the fit(), set_parameters() and package().

You can find all details in FedAvgClient and FedAvgServer, which are the bases of all implementations in FL-bench.

Integrating Dataset

• Inherit your own dataset class from BaseDataset in data/utils/datasets.py and add your class in dict DATASETS. Highly recommend to refer to the existing dataset classes for guidance.

Customizing Model

• I offer the CustomModel class in src/utils/models.py and you just need to define your model arch.
• If you want to use your customized model within FL-bench's workflow, the base and classifier must be defined. (Tips: You can define one of them as torch.nn.Identity() for bypassing it.)

Citation 🧐

@software{Tan_FL-bench,
  author = {Tan, Jiahao and Wang, Xinpeng},
  license = {MIT},
  title = {{FL-bench: A federated learning benchmark for solving image classification tasks}},
  url = {https://github.com/KarhouTam/FL-bench}
}

@misc{tan2023pfedsim,
  title={pFedSim: Similarity-Aware Model Aggregation Towards Personalized Federated Learning}, 
  author={Jiahao Tan and Yipeng Zhou and Gang Liu and Jessie Hui Wang and Shui Yu},
  year={2023},
  eprint={2305.15706},
  archivePrefix={arXiv},
  primaryClass={cs.LG}
}