Skip to content

Liangqiong/ViT-FL-main

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

19 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Vision Transformer in Federated Learning

Usage

0. Installation

  • Run cd ViT-FL-main
  • Install the libraries listed in requirements.txt

1. Prepare Dataset

We provide the data partitions for Cifar-10 and CelebA datasets

2. Set (download) the Pretrained Models

3. Train Model (refer to additional notes for usage of more tags)

  • ViT-CWT train on Split-2, KS-0.65 of Cifar-10 and real world federated dataset CelebA
python train_CWT.py --FL_platform ViT-CWT --net_name ViT-small --dataset cifar10 --E_epoch 1 --max_communication_rounds 100 --split_type split_2 --save_model_flag
python train_CWT.py --FL_platform ViT-CWT --net_name ViT-small --dataset CelebA --E_epoch 1 --max_communication_rounds 30 --split_type real

  • ViT-FedAVG train on Split-2, KS-0.65 of Cifar-10 and real world federated dataset CelebA
python train_FedAVG.py --FL_platform ViT-FedAVG --net_name ViT-small --dataset cifar10 --E_epoch 1 --max_communication_rounds 100 --num_local_clients -1 --split_type split_2 --save_model_flag
python train_FedAVG.py --FL_platform ViT-FedAVG --net_name ViT-small --dataset CelebA --E_epoch 1 --max_communication_rounds 30 --num_local_clients 10 --split_type real

  • All the checkpoints, results, log files will be saved to the --output_dir folder, with the final performance saved at log_file.txt

Additional Notes

  • Some important tags for both train_CWT.py and train_FedAVG.py:

    • --FL_platform: selection of FL platforms, ViT-CWT, ResNet-CWT, EfficientNet-CWT, or Swin-CWT for train_CWT.py, ViT-FedAVG, ResNet-FedAVG, EfficientNet-FedAVG, or Swin-FedAVG for train_FedAVG.py
    • --net_name: basic Name of this run, also providing detailed network-architecture for ViT/ResNet/EfficientNet. For ViT: ViT-small, ViT-tiny, ViT-base(default), For EfficientNet: efficientnet-b1, efficientnet-b5(default), efficientnet-b7 see sstart_config.py for more details
    • --dataset: choose of the following three datasets ["cifar10", "Retina" ,"CelebA"]
    • --save_model_flag: set to True if need to save the checkpoints
    • --output_dir: the output directory where checkpoints/results/logs will be written
    • --decay_type: learning rate decay schedulers with the following three options ["cosine", "linear", "step"]
    • --E_epoch: local training epoch E in FL train
    • --max_communication_rounds: total communication rounds, 100 for Retina and Cifar-10, 30 for CelebA
    • --split_type: type of data partitions, supports ["split_1", "split_2", "split_3"] for Cifar-10 and Retina, ["real"] for CelebA
    • --cfg: configuration document for Swin-transformers if use Swin-FL, otherwise ignored it
  • Additional tag for paralle FedAVG

    • --num_local_clients: Num of local clients joined in each FL train. -1 (usage of all local clients) for Retina and Cifar-10, 10 for CelebA.
  • Also refer to the train_CWT.py and train_FedAVG.py for more tags

Acknowledgments

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages