DOCSEARCHbyAlgolia
Submit Docs

PFLlib: Personalized Federated Learning Library and Benchmark

repository·master·Indexed 24 days ago

https://github.com/tsingz0/pfllib

A beginner-friendly library and benchmark for Personalized Federated Learning (pFL) and traditional Federated Learning (tFL). It supports 39 algorithms (including FedAvg, SCAFFOLD, FedRep, and FedDBE), 24 datasets, and 3 data heterogeneity scenarios: label skew, feature shift, and real-world. PFLlib provides tools for privacy evaluation using DLG and PSNR, simulation of network constraints like client dropout and slow trainer rates, and an extensible framework for adding new algorithms, models, and datasets.

Tokens
3.6K
Snippets
5
Records
16
Agent score
30%

What's inside PFLlib

  1. Overview of PFLlib

    master

    PFLlib is a beginner-friendly Personalized Federated Learning (pFL) library and benchmark designed to help users master Federated Learning (FL) quickly. It supports 39 traditional FL (tFL) and personalized FL (pFL) algorithms, 3 scenarios, and 24 datasets. The platform is optimized for efficiency, capable of simulating 500 clients using a 4-layer CNN on Cifar100 with only ~5.08GB of GPU memory on a single NVIDIA GeForce RTX 3090.

    Key capabilities include:

    • Algorithm Support: A wide range of tFL (e.g., FedAvg, SCAFFOLD, FedProx) and pFL (e.g., Per-FedAvg, FedRep, FedDBE) algorithms.
    • Scenario Simulation: Support for various data heterogeneity scenarios like label skew and feature shift.
    • Evaluation: Provides privacy evaluation and systematic research support.
    • New Client Testing: Ability to train on a subset of clients and evaluate performance on new, unseen clients by configuring args.num_new_clients in ./system/main.py (note: not all algorithms support this).

  2. Understand the supported Federated Learning scenarios

    master

    PFLlib supports three primary data distribution scenarios for testing Personalized Federated Learning:

    1. label skew: Focuses on non-IID data where labels are distributed unevenly. It includes two sub-types:
      • Pathological non-IID: Each client holds only a small subset of available labels (e.g., 2 out of 10).
      • Practical non-IID: Uses a Dirichlet distribution to model more realistic, less extreme imbalances.
      • Note: You can use the balance option to ensure data amounts are evenly distributed across clients.
    2. feature shift: Focuses on Domain Adaptation using datasets like Amazon Review, Digit5, and DomainNet.
    3. real-world: Uses naturally separated datasets such as Camelyon17 (hospitals), iWildCam (camera traps), Omniglot, HAR, and PAMAP2.

  3. How to extend PFLlib with new algorithms and data

    master

    PFLlib is designed to be easily extensible. To add new components, follow these patterns:

    Adding a new algorithm

    To implement a new algorithm, you need to add new features to the client and server logic files:

    1. Modify clientNAME.py to define client-side behavior.
    2. Modify serverNAME.py to define server-side aggregation or coordination.

    Adding new data

    Use the generate_DATA.py script to create new scenarios. You can then run these scenarios using the existing framework by invoking main.py, clientNAME.py, and serverNAME.py.

  4. Add a new model or optimizer

    master

    To extend the library's capabilities with new architectures or optimization strategies:

    • New Model: Add the model definition to ./system/flcore/trainmodel/models.py.
    • New Optimizer: Add the optimizer implementation to ./system/flcore/optimizers/fedoptimizer.py.

  5. Set up the PFLlib environment

    master

    PFLlib requires CUDA and Conda.

    1. Install CUDA.
    2. Install conda and activate it.
    3. Create the environment using the provided YAML file. You may need to downgrade torch via pip if your CUDA version requires it.

    For additional configurations, refer to the prepare.sh script in the repository.

    conda env create -f env_cuda_latest.yaml

  6. Add a new algorithm to PFLlib

    master

    New algorithms are implemented by extending the base Server and Client classes.

    Implementing a Server

    Extend Server from ./system/flcore/servers/serverbase.py. You must implement the train method to define your algorithm's server-side scheduling logic.

    Implementing a Client

    Extend Client from ./system/flcore/clients/clientbase.py. You must implement the train method to define the client-side training logic.

    Refer to the following patterns for implementation:

    # serverNAME.py
import necessary pkgs
from flcore.clients.clientNAME import clientNAME
from flcore.servers.serverbase import Server

class NAME(Server):
    def __init__(self, args, times):
        super().__init__(args, times)

        # select slow clients
        self.set_slow_clients()
        self.set_clients(clientAVG)
    def train(self):
        # server scheduling code of your algorithm
    # clientNAME.py
import necessary pkgs
from flcore.clients.clientbase import Client

class clientNAME(Client):
    def __init__(self, args, id, train_samples, test_samples, **kwargs):
        super().__init__(args, id, train_samples, test_samples, **kwargs)
        # add specific initialization
    
    def train(self):
        # client training code of your algorithm

  7. Add a new dataset to PFLlib

    master

    To add a new dataset, create a generate_DATA.py file in the ./dataset directory. Use ./dataset/generate_MNIST.py as a template. Your script should handle downloading, preprocessing, and splitting the data using the provided utilities.

    Follow this structure:

    1. Import necessary packages and processing functions from utils.
    2. Define a generate_dataset function.
    3. Use separate_data to split content from labels.
    4. Use split_data to create training and testing sets.
    5. Use save_file to persist the processed data to the configured paths.
    # `generate_DATA.py`
import necessary pkgs
from utils import necessary processing funcs

def generate_dataset(...):
  # download dataset as usual
  # pre-process dataset as usual
  X, y, statistic = separate_data((dataset_content, dataset_label), ...)
  train_data, test_data = split_data(X, y)
  save_file(config_path, train_path, test_path, train_data, test_data, statistic, ...)

# call the generate_dataset func

  8. Run Federated Learning evaluations

    master

    After cloning the repository, setting up the environment, and generating your datasets, you can run evaluations using main.py inside the ./system directory.

    Command Arguments:

    • -data: The dataset to use (e.g., MNIST).
    • -m: The model architecture (e.g., CNN).
    • -algo: The federated learning algorithm (e.g., FedAvg).
    • -gr: Number of rounds (e.g., 2000).
    • -did: Device ID (e.g., 0 for single GPU, or a comma-separated list like 0,1,2,3 for multiple GPUs).

    Note: It is recommended to tune algorithm-specific hyperparameters before running on new hardware.

    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
python main.py -data MNIST -m CNN -algo FedAvg -gr 2000 -did 0,1,2,3 # running on multiple GPUs

  9. Evaluate privacy using DLG and PSNR

    master

    PFLlib supports privacy evaluation to assess the security of tFL/pFL algorithms.

    • Supported Attack: DLG (Deep Leakage from Gradients).
    • Supported Metric: PSNR (Peak Signal-to-Noise Ratio). For image evaluation, a lower PSNR score indicates better privacy-preserving capabilities (higher leakage).

    For a concrete implementation example, refer to ./system/flcore/servers/serveravg.py.

  10. Generate MNIST datasets for label skew scenarios

    master

    To generate datasets for MNIST under different label skew conditions, navigate to the ./dataset directory and use the generate_MNIST.py script.

    Important: Before running, you may need to modify train_ratio and alpha in dataset/utils/dataset_utils.py to customize the distribution.

    Available command patterns:

    • IID (Unbalanced): python generate_MNIST.py iid - -
    • IID (Balanced): python generate_MNIST.py iid balance -
    • Pathological non-IID (Unbalanced): python generate_MNIST.py noniid - pat
    • Practical non-IID (Unbalanced): python generate_MNIST.py noniid - dir
    • Extended Dirichlet strategy: python generate_MNIST.py noniid - exdir
    cd ./dataset
# Please modify train_ratio and alpha in dataset\utils\dataset_utils.py

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

  11. Evaluate performance on new clients

    master

    You can simulate a scenario where the model is trained on a specific set of clients and then evaluated on a separate set of new clients. This is controlled via the args.num_new_clients argument in the ./system/main.py script.

    Note: This feature is not supported by all tFL/pFL algorithms implemented in the library.

  12. Run federated learning experiments via system/main.py

    master

    The system/main.py script is the primary entrypoint for orchestrating federated learning experiments. It handles model initialization, algorithm selection, and the training loop. You run it from the command line by passing various arguments to configure the dataset, model architecture, federated algorithm, and hyper-parameters.

    To run an experiment, you typically use the Python interpreter to call the script with desired flags. The script automatically handles device assignment (CPU/CUDA) and manages the lifecycle of the server and clients.