Home

Awesome

Federated Learning with Parial Model Personalization

Authors: Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael Rabbat, Maziar Sanjabi, Lin Xiao

Contact: pillutla@cs.washington.edu, linx@fb.com

Paper Link.

Requirements

Install PyTorch version 1.9 for your appropriate cuda version, as well torchvision and torchaudio. Then, install other dependences with pip install -r requirements.txt.

Outline of the Repository

Here is the outline:

The overall pipeline for each task is:

  1. Pretrain the model without personalization (i.e., using FedAvg variants)
  2. Personalize with partial model personalization (using the proposed FedAlt or FedSim) or pFedMe. Note that this stage does not apply for finetuning or Ditto (finetuning with L2 regularization)
  3. Finetune locally (for PFL, pFedMe and Ditto)

Main files

The two main files are:

Scripts

Here is the breakdown of the directory scripts/federated. This contains the scripts to run steps 1 and 2 of the pipeline.

Here is the breakdown of the directory scripts/finetune. This containts the scripts to run step 3 of the pipeline.

Example use of scripts

The scripts for the experimental pipeline are:

  1. Pretrain in a non-personalized manner with FedAvg: scripts/federated/pretrain/stackoverflow.sh
  2. Run FedAlt: scripts/federated/partial_personalization/stackoverflow.sh - Note: the saved model in step 1 is loaded again in step 2. Make sure the argument pretrained_model_path points to the correct checkpoint from step 1
  3. Local Finetuning: scripts/finetune/partial_personalization/stackoverflow.sh

Implementation details

The codebase is constructed in a modular and extensible manner. The main functionality is divided into 3 parts: data, models and optimization. Correspondingly, the folder pfl contains the following subfolders:

Each of these See the README in each of these folders for more details. The details on how to add new datasets or new optimization algorithms are given there. In the current implementation, a new dataset also requires models specific to these datasets to be implemented.

Future work to improve the code

License

This project is released under MIT License, which allows commercial use. See LICENSE for details.

Citation

If you find this repository useful, please consider giving a star :star: and citation:

@inproceedings{Pillutla2022pfl,
  author  = {Krishna Pillutla, Kshitiz Malik, Abdelrahman Mohamed, Michael Rabbat, Maziar Sanjabi, Lin Xiao},
  title   = {Federated Learning with Partial Model Personalization},
  booktitle = {Proceedings of the 39th International Conference on Machine Learning (ICML'22)},
  year    = {2022},
}