FedFound: Federated Foundation Model for GI Endoscopy Imaging

This repository contains the PyTorch implementation of FedFound, our federated learning framework for foundation model pretraining and downstream tasks in gastrointestinal (GI) endoscopy imaging.

FedFound enables federated self-supervised pretraining of Vision Transformer (ViT)–based foundation models across multiple institutions while preserving data privacy. The learned representations can be fine-tuned for classification, object detection, and instance segmentation in GI endoscopy datasets.

Quick Start

1. Create Conda environment

conda create --name ffm python=3.8

2. Install required libraries

pip install -r requirements.txt

3. Dataset Preparation

Please follow the instructions in data/README.md to organize your dataset before training.

4. Run pretraining

The scripts for pretraining and fine-tuning are located in the /code/fedmae/scripts/endo folder:

• SLURM users:
  Run start_multiple_pretraining.sh to launch federated pretraining across multiple nodes.

• Local users:
  Run individual pretraining scripts such as fedavg_ssl.sh, fedavgm_ssl.sh, etc. for single-machine experiments.

• Baseline experiments:
  Use fed_ssl.sh to run lower- and upper-bound comparisons.

5. Finetuning

The scripts for fine-tuning pretrained models are also in the /code/fed_mae/scripts/endo folder:

• SLURM users:
  Run start_multiple_finetuning.sh to launch fine-tuning for multiple methods.

• Local users:
  Refer to fed_ssl_finetune_multiple.sh for individual or single-machine fine-tuning experiments.

Acknowledgements

This repository builds upon and extends the code from the SSL-FL repository by Rui Yan et al., which provided the starting point for our federated self-supervised learning framework.