Overview

Implementation of FedAvg in Pytorch, supporting WandB logging.

Installation

Dependencies

Requires Python 3.6+.

• numpy>=1.22.4
• pytorch>=1.12.0
• torchvision>=0.13.0
• wandb>=0.12.19

Conda Installation

conda env create -f environment.yml

How To Use

Major Arguments

Flag	Options	Default	Info
--data_root	String	"../datasets/"	path to data directory
--model_name	String	"cnn"	name of the model (cnn, mlp)
--non_iid	Int (0 or 1)	1	0: IID, 1: Non-IID
--n_clients	Int	100	number of the clients
--n_shards	Int	200	number of shards
--frac	Float	0.1	fraction of clients in each round
--n_epochs	Int	1000	total number of rounds
--n_client_epochs	Int	5	number of local training epochs
--batch_size	Int	10	batch size
--lr	Float	0.01	leanring-rate
--wandb	Bool	False	log the results to WandB

Training Example

python fed_avg.py --batch_size=10 --frac=0.1 --lr=0.01 --n_client_epochs=20 --n_clients=100 --n_epochs=1000 --n_shards=200 --non_iid=1

To perform a sweep over hyperparameters using WandB:

wandb sweep sweep.yaml
wandb agent <sweep_id>