SplitAVG

This repository includes the official project of SplitAVG, from our paper "SplitAVG: A heterogeneity-aware federated deep learning method for medical imaging"(https://ieeexplore.ieee.org/abstract/document/9806163?casa_token=AiCs74_RK4EAAAAA:fw3aflyOBRq3_E9y7rgkec4bCC93UWUEJF-OtjmUgCheQONFhkNhUdNh8ObPMB9_5fkeHKCXOQ), accepted to IEEE Journal of Biomedical and Health Informatics (JBHI) 2022.

Data preparation

Our simulated federated data partitions for Kaggle bone age regression task can be downloaded at: https://drive.google.com/drive/folders/1O8oJ2KBHUprvRo0ILD93XfiYaW8u0Q9I?usp=sharing (Custom data could be used with the same format as below)

The training data consists of 4 hetergeneous splits (institutions). The organization structures of the training data and validation data are:

boneS1.h5
  -- ['examples']  -- ['0']   -- ['pixels']
                              -- ['label']
                   -- ['1']   -- ['pixels']
                              -- ['label']    
                   -- ['2']   -- ['pixels']
                              -- ['label'] 
                   -- ['3']   -- ['pixels']
                              -- ['label']

val.h5
  -- ['examples']  -- ['0']   -- ['pixels']
                              -- ['label']

Visualization of the heterogeneity in label distributions across institutions of bone age dataset:

Train

python main_SplitAVG.py \
--batch_size 32 \
--train_file "./data/boneS1.h5" --val_file "./data/val.h5" --site_num 4 --iter_per_epoch 71 --num_class 1 \
--arch 'res34' \
--splitavg_v2 False \
--seed 2556 \