Enhancing Clustered Federated Learning: Integration of Strategies and Improved Methodologies

If our project helps you, please give us a star ⭐ and cite our paper!

Environment Setup

To run the experiments on the DelftBlue HPC (or locally), follow these setup steps:

1. Conda Environment Setup

We have provided an environment.yml and a helper script setup_env.sh to initialize the environment with all the correct package versions (PyTorch 2.2 with CUDA 12.1 support, OpenCV, SciPy, cvxpy, etc.).

If you are on macOS, we have provided environment_mac.yml which removes the CUDA-specific dependencies.

On DelftBlue (or Linux with CUDA), run the setup script:

bash setup_env.sh

Or manually create the environment using Conda:

For DelftBlue / Linux (CUDA):

# Load required modules on DelftBlue
module load 2025
module load cuda/12.9
module load miniconda3

# Create and activate environment
conda env create -f environment.yml
conda activate hcfl

For macOS:

# Create and activate environment
conda env create -f environment_mac.yml
conda activate hcfl

2. Cache Pretrained Models (For Compute Nodes without Internet)

Since DelftBlue compute/GPU nodes do not have internet access, they cannot download pretrained model weights (like MobileNet-V2 or ResNet-18) at runtime.

To download and cache the necessary weights prior to running your training jobs, run the following commands on the login node (which has internet access):

# Activate the environment
conda activate hcfl

# Download and cache MobileNet-V2 (for CIFAR-100 runs)
python -c "import torchvision.models as models; models.mobilenet_v2(pretrained=True)"

# Download and cache ResNet-18 (for CIFAR-10 runs)
python -c "import torchvision.models as models; models.resnet18(pretrained=True)"

These weights will be saved to your user's global cache directory ~/.cache/torch/hub/checkpoints/ and will automatically be loaded by compute nodes during training.

3. Dataset Creation

Before running any experiments, you need to construct the federated datasets.

# Create the necessary directory structure
mkdir -p logs/slurm
mkdir -p data/cifar100-c-2swap
mkdir -p data/cifar100-c/origin

# Generate the federated dataset
python create_c/make_cifar100_c-2swap.py

This script downloads CIFAR-100, splits it, injects the concept-shift swaps, and populates the data/cifar100-c-2swap/ directory with pickle partitions.

Requirements:

The complete environment specification is in environment.yml. The original paper's version requirements are listed below for reference:

Pillow == 8.1.2
tqdm
scikit-learn == 0.21.3
numpy == 1.19.0
torch == 1.2.0
matplotlib == 3.1.1
networkx == 2.5.1
cvxpy
torchvision
tensorboard

Configurations:

• Aggregator: Please refer to FedIASAggregator class of aggregator.py, following the annotations to choose the backbone algorithms, aggregation methods, and similarity metrics.
• Client similarity: Please refer to SplitLearnersEnsemble class of learners\learners_ensemble.py, following the annotations to choose using graidents or prototypes to calculate the similarity.

Example scripts:

python run_experiment.py cifar10-c-2swap FedEM_SW --n_learners 2 --n_rounds 200 --bz 128 --lr 0.03 --lr_scheduler constant --log_freq 1 --device 1 --optimizer sgd --seed 1 --verbose 1 --suffix 03-lr-03-resnet-split-005-04-1 --split

Hints:

• All the supported algorithms can be found in constants.py. The default models are avaliable in models.py. The detailed explanation of arguments can be found in args.py.
• You need to first use scrips in create_c to construct the datasets.
• The --split indicates using shared feature extractor for all clients.

Acknowledgement

We are grateful for the following awesome projects:

• FedEM
• Robustness Benchmark
• FedSoft
• IFCA
• FedRC

Bibliography

If you find this repository helpful for your project, please consider citing:

@inproceedings{
guo2025enhancing,
title={Enhancing Clustered Federated Learning: Integration of Strategies and Improved Methodologies},
author={Yongxin Guo and Xiaoying Tang and Tao Lin},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=zPDpdk3V8L}
}