FedZoo-Bench, is an open-source library based on PyTorch that facilitates experimentation in federated learning by providing researchers with a comprehensive set of standardized and customizable features such as training, Non-IID data partitioning, fine-tuning, performance evaluation, fairness assessment, and generalization to newcomers, for both global and personalized FL approaches.
Most modules are put in the src directory. The main modules to run the code (which you may modify to customize your experiments) are organized as follows:
| Directory | Contents |
|---|---|
src/benchmarks |
Server modules and main training functions for each algorithm. |
src/client |
Client modules. |
src/data |
Dataset pre-processing and loaders. |
src/models |
Neural network architectures. |
src/utils |
Other utils and argument parsers. |
9 Global FL algorithms and 15 Personalized FL algorithms are supported in this code.
Global Algorithms:
| Algorithm | Publication | Code |
|---|---|---|
| FedAvg | PMLR, 2017, https://arxiv.org/abs/1602.05629 | |
| FedProx | MLSys, 2020, https://arxiv.org/abs/1812.06127 | https://github.com/litian96/FedProx |
| FedNova | NeurIPS, 2020, https://arxiv.org/abs/2007.07481 | https://github.com/JYWa/FedNova |
| Scaffold | PMLR, 2020, https://arxiv.org/abs/1910.06378 | |
| MOON | CVPR, 2021, https://arxiv.org/abs/2103.16257 | https://github.com/QinbinLi/MOON |
| FedBN | ICLR, 2021, https://arxiv.org/abs/2102.07623 | https://github.com/med-air/FedBN |
| FedDyn | ICLR, 2021, https://arxiv.org/abs/2111.04263 | https://github.com/alpemreacar/FedDyn |
| FedDF | NeurIPS, 2020, https://arxiv.org/abs/2006.07242 | |
| FedAvgM | https://arxiv.org/abs/1909.06335 | |
| FedavgM | https://arxiv.org/abs/2002.06440 |
Personalized Algorithms:
We support 9 widely used datasets: MNIST, CIFAR-10, CFIAR-100, USPS, SVHN, FMNIST, FEMNIST, Tiny-ImageNet, and STL-10.
The datasets should be aligned as follows:
datadir
├── cifar-10-batches-py
│ └── ...
├── cifar-100-python
│ └── ...
├── fmnist
│ └── ...
...
and the argument datadir should point to this datadir.
Most dataests will be downloaded automatically if they are not in the directory except Tiny-ImageNet, which is available at http://cs231n.stanford.edu/tiny-imagenet-200.zip.
We provide scripts to run the algorithms, which are put under scripts/. Here is an example to run the script:
cd scripts
bash fedavg.sh
Custom experiments can be performed by running main.py with customized arguments. For example:
python main.py --ntrials 3 ...
The descriptions of general arguments are as follows:
| Parameter | Description |
|---|---|
| ntrials | The number of total runs. |
| rounds | The number of communication rounds per run. |
| num_users | The number of clients. |
| nclass | Classes or shards per user. |
| nsample_pc | The number of samples per class or shard for each client. |
| frac | The fraction of clients updated per round. |
| local_ep | The number of local training epochs. |
| local_bs | Local batch size. |
| lr | The learning rate for local models. |
| momentum | The momentum for the optimizer. |
| model | Network architecture. Options: lenet5, simple-cnn-3, resnet9, resnet, vgg16. |
| dataset | The dataset for training and testing. Options are discussed above. |
| partition | How datasets are partitioned. Options: homo, noniid-labeldir, noniid-#label1 (or 2, 3, ..., which means the fixed number of labels each party owns). |
| datadir | The path of datasets. |
| logdir | The path to store logs. |
| log_filename | The folder name for multiple runs. E.g., with ntrials=3 and log_filename=$trial, the logs of 3 runs will be located in 3 folders named 1, 2, and 3. |
| alg | Federated learning algorithm. Options are discussed above. |
| niid_beta | The parameter for non-iid data partitioning. |
| iid_beta | The parameter for iid data partitioning. |
| gpu | The IDs of GPU to use. E.g., 0. |
| print_freq | The frequency to print training logs. E.g., with print_freq=10, training logs are displayed every 10 communication rounds. |
Algorithm specific arguments:
| Parameter | Description |
|---|---|
| local_rep_ep | The number of local rep layers updates for FedRep. |
| feddf_n | The number of server-side model fusion: N. |
| glob_layers | The number of global or personal layers. |
| Parameter | Description |
|---|---|
| mu | FedProx regularizer parameter. default=0.001. (float) |
| Parameter | Description |
|---|---|
| lam_ditto | Ditto parameter lambda. default=0.8. (float) |
| Parameter | Description |
|---|---|
| alpha_apfl | APFL parameter Alpha. default=0.75. (float) |
| Parameter | Description |
|---|---|
| use_project_head | whether use projection head or not. (action store true) |
| temperature_moon | Moon parameter Temperature. default=0.5. (float) |
| mu_moon | Moon parameter Mu. default=1.0. (float) |
| Parameter | Description |
|---|---|
| pacfl_beta | PACFL clustering threshold. |
| pacfl_n_basis | PACFL number of basis per label. |
| pacfl_linkage | PACFL Type of Linkage for HC. |
| nclasses | The mumber of classes for PACFL. |
| Parameter | Description |
|---|---|
| flis_cluster_alpha | FLIS clustering threshold. default=0.5. (float) |
| nclasses | The mumber of classes for public dataset. default=10. (int) |
| Parameter | Description |
|---|---|
| pfedme_beta | pFedMe beta for global model update. default=1.0. (float) |
| Parameter | Description |
|---|---|
| nclusters | The number of Clusters for IFCA. default=2. (int) |
| Parameter | Description |
|---|---|
| pruning_percent_subfedavg | Pruning percent for layers (0-100) for subfedavg. default= 5. (float) |
| pruning_target_subfedavg | Total Pruning target percentage (0-100) for subfedavg. default= 35. (float) |
| dist_thresh_subfedavg | The threshold for fcs masks difference for subfedavg . default=0.0001. (float) |
| acc_thresh_subfedavg | The accuracy threshold to apply the derived pruning mask for subfedavg. default=55. (float) |
| Parameter | Description |
|---|---|
| n_models | The number of mixture distributions M for FedEM. |
For generalization to new comers experiment, please use the following arguments in your script.
| Parameter | Description |
|---|---|
| new_comer | If passed in, 80% clients are involved for FL training and 20% are evaluated afterwards as new comers. |
| ft_epoch | The number of finetune epochs for new comers. |
If you find this repository useful, please cite our work:
For any issues please feel free to submit an issue. We are also open to any contributions.