FKD-Med: Privacy-Aware, Communication-Optimized Medical Image Segmentation via Federated Learning and Model Lightweighting through Knowledge Distillation
This project is licensed under the GNU General Public License v3.0 - see the LICENSE file for details.
This repository contains code for training and testing U-Net for image semantic segmentation tasks. It contains both traditional and Federated Traning using the FedAvg algorithm in the Flower framework. All operations of the user are done in this directory
https://paperswithcode.com/dataset/cvc-clinicdb
https://www.kaggle.com/datasets/nikhilpandey360/chest-xray-masks-and-labels
In the /model/Unet_model
,training on different Unet models,
In the Loss.py
,choosing different loss functions for training
In the DataSet.py
,selecting different medical data for training
The model is customized in the command line for the reader to choose from
python3 train.py --client client1 --cuda cuda1 --model resUnet --num_epochs 50 --dataset CVC --picFormat .tif
python3 train.py --client client2 --cuda cuda2 --model resUnet --num_epochs 50 --dataset CVC --picFormat .tif
python3 train.py --client client3 --cuda cuda3 --model resUnet --num_epochs 50 --dataset CVC --picFormat .tif
The model is customized in the command line for the reader to choose from
python3 train.py --client client1 --cuda cuda1 --model resUnet --num_epochs 50 --dataset Chest --picFormat .png
python3 train.py --client client2 --cuda cuda2 --model resUnet --num_epochs 50 --dataset Chest --picFormat .png
python3 train.py --client client3 --cuda cuda3 --model resUnet --num_epochs 50 --dataset Chest --picFormat .png
python3 server.py
python3 client.py --client client1 --cuda cuda1 --model resUnet --num_epochs 10 --dataset CVC --picFormat .tif
python3 client.py --client client2 --cuda cuda2 --model resUnet --num_epochs 10 --dataset CVC --picFormat .tif
python3 client.py --client client3 --cuda cuda3 --model resUnet --num_epochs 10 --dataset CVC --picFormat .tif
python3 server.py
python3 client.py --client client1 --cuda cuda1 --model resUnet --num_epochs 10 --dataset Chest --picFormat .png
python3 client.py --client client2 --cuda cuda2 --model resUnet --num_epochs 10 --dataset Chest --picFormat .png
python3 client.py --client client3 --cuda cuda3 --model resUnet --num_epochs 10 --dataset Chest --picFormat .png
The teacher model has been trained in the train.py
file according to client,line 202,weights are found inside saved_models/
folder Eg.
torch.save(model.state_dict(),PATH)
The trained teacher model was used at line 129 of clientFKD.py
teacher.load_state_dict(torch.load(PATH))
python3 server.py
python3 clientFKD.py --client client1 --cuda cuda1 --model resUnet --num_epochs 10 --dataset CVC --picFormat .tif
python3 clientFKD.py --client client2 --cuda cuda2 --model resUnet --num_epochs 10 --dataset CVC --picFormat .tif
python3 clientFKD.py --client client3 --cuda cuda3 --model resUnet --num_epochs 10 --dataset CVC --picFormat .tif
python3 server.py
python3 clientFKD.py --client client1 --cuda cuda1 --model resUnet --num_epochs 10 --dataset Chest --picFormat .png
python3 clientFKD.py --client client2 --cuda cuda2 --model resUnet --num_epochs 10 --dataset Chest --picFormat .png
python3 clientFKD.py --client client3 --cuda cuda3 --model resUnet --num_epochs 10 --dataset Chest --picFormat .png