FedMef

This repository provides an implementation for paper "FedMef: Towards Memory-efficient Federated Dynamic Pruning" and the unofficial implementation for FedTiny and FedDST.

Getting Started

Set the environment

conda create -n fedmef python=3.8
conda activate fedmef
conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=10.2 -c pytorch
conda install -c anaconda mpi4py
pip install -r requirements.txt

Download the dataset

cd data/cifar10/
bash download_cifar10.sh
cd ../..

Setting ip configurations for grpc

1. create .csv file in the format:

    receiver_id,ip
    0,<ip_0>
    ...
    n,<ip_n>
    
    where n = client_num_per_round

2. provide path to file as argument to --grpc_ipconfig_path

Running using TRPC

In order to run using TRPC set master's address and port in file trpc_master_config.csv, and use TRPC as backend option.

Set GPU and process number

Modify the gpu_mapping.yaml to specify the process, e.g., for 4 GPUs

mapping_default:
    Server: [3, 3, 3, 2]

Run the code

The command to run the code

run.sh 
    CLIENT_NUM=$1 # the total number of clients
    WORKER_NUM=$2 # the number of selected clients
    MODEL=$3 # base model 
    ROUND=$4 # the number of communication rounds
    EPOCH=$5 # the number of local training epoch
    DATASET=$6 # the name of dataset
    DATA_DIR=$7 # the dir of dataset
    PRUNE=$8 # the type of pruning methods
    DENSITY=$9 # the density of model
    BaE=${10} # Use Budget-aware Extrusion (BaE) 
    SAP=${11} # Use Scaled Activation Pruning (SAP) 
    BATCH=${12} # Batch Size

For example:

bash run.sh 100 10 resnet18 500 10 cifar10 "data/cifar10" FedMef 0.1 1 1 64 
bash run.sh 100 10 resnet18 500 10 cifar10 "data/cifar10" FedTiny 0.1 0 0 64  
bash run.sh 100 10 resnet18 500 10 cifar10 "data/cifar10" FedDST 0.1 0 0 64 
bash run.sh 100 10 resnet18 500 10 cifar10 "data/cifar10" Mag 0.1 0 0 64 
bash run.sh 100 10 resnet18 500 10 cifar10 "data/cifar10" None 1.0 0 0 64