AdwU-Net: Adaptive depth and Width U-Net
Built upon MIC-DKFZ/nnUNet, this repository provides the official PyTorch implementation of AdwU-Net: Adaptive Depth and Width U-Net for Medical Image Segmentation by Differentiable Neural Architecture Search.
Environments and Requirements:
1.Install nnU-Net as below. You should meet the requirements of nnUNet, our method does not need any additional requirements. For more details, please refer to https://github.com/MIC-DKFZ/nnUNet
git clone https://github.com/MIC-DKFZ/nnUNet.git
cd nnUNet
pip install -e .
2.Set environment variables for nnU-Net. Concretely, Set the paths in your .bashrc file, which is located in your home directory. Open the file and add the following lines to the bottom:
export nnUNet_raw_data_base="/data/hzy/nnUNet/nnUNet_raw"
export nnUNet_preprocessed="/data/hzy/nnUNet/nnUNet_preprocessed"
export RESULTS_FOLDER="/data/hzy/nnUNet/nnUNet_trained_models"
(of course adapt the paths to your system)
3.Copy the python files in this repository to the code directory of nnUNet.
cp network_training/* nnunet/training/network_training/
cp network_architecture/* nnunet/network_architecture/
cp run_searching.py nnunet/run/
Dataset and Preprocessing
The training data is from the MSD (Medical Segmentation Decathlon) challenge.
After downloading the data, use scripts nnUNet_convert_decathlon_task and nnUNet_plan_and_preprocess provided by nnUNet to preprocess the data.
For more details, please refer to the repository of nnUNet.
How to use AdwU-Net
For example, we search the optimal depth and width of Task005_Prostate and then configure the optimal depth and width of nnUNet to train the model.
- Run the following command to search the optimal depth and width on Task005_Prostate
cd nnunet/run
python run_searching.py 3d_fullres AdwUNetTrainer_search 5 all
- Edit AdwUNetTrainer.py, change the value of self.arch_list using the output at the end of search stage.
- Run the following command to train nnUNet using the optimal depth and width on Task005_Prostate
for FOLD in 0 1 2 3 4
do
nnUNet_train 3d_fullres AdwUNetTrainer 5 $FOLD
done
Citation
If you find this repository useful, please consider citing our paper:
@inproceedings{
huang2022adwunet,
title={AdwU-Net: Adaptive Depth and Width U-Net for Medical Image Segmentation by Differentiable Neural Architecture Search},
author={Ziyan Huang and Zehua Wang and zhikai yang and Lixu Gu},
booktitle={Medical Imaging with Deep Learning},
year={2022},
url={https://openreview.net/forum?id=kF-d1SKWJpS}
}