nnunet代码使用

nnunet作为一个医学分割中神一样的代码。有很多地方可能在运行时需要改进。这就要求我们不能仅满足于安装nnunet包，并通过官方命令运行。而应对代码进行修改。

本文主要讲解在下载nnunet后如何通过修改代码运行。

关于nnunet的讲解，和安装nnunet包，使用命令运行的方法推荐这个博主的四篇文章。

 

阅读本文是可以同时参考官方使用教程

1、设置conda环境

为了避免环境适配的问题，最好安装一个新的conda环境。python3, torch1.6以上。

温馨提供：安装conda环境教程

2、下载nnunet

git clone https://github.com/MIC-DKFZ/nnUNet.git

3、修改使用nnunet

此处首先说明一点。由于运行nnunet中经常调用上级包甚至上级的上级；比如最常见的错误是找不到nnunet,"no module named 'nnunet'"。因此可使用：

import sys
sys.path.append("/home2/*/nnUNet/")    # 添加nnUNet的绝对目录到环境变量中

3.1 整理初始数据

由于nnunet是需要固定输入格式的，具体格式详见网站

因为为了制作这个格式的数据，可以使用nnUNet/nnunet/dataset_conversion/ 中的代码进行操作。此处以Task029_XXX.py 为例。

train_dir = "/media/fabian/DeepLearningData/tmp/LITS-Challenge-Train-Data"
test_dir = "/media/fabian/My Book/datasets/LiTS/test_data"

output_folder = "/media/fabian/MyBook/MedicalDecathlon/MedicalDecathlon_raw_splitted/Task029_LITS"

...
nii_files_tr_data = subfiles(train_dir, True, "volume", "nii", True)   # 获得train_dir中包含volume,且以nii结尾的数据
nii_files_tr_seg = subfiles(train_dir, True, "segmen", "nii", True)    # 获得train_dir中包含segmen，且以nii结尾的数据
 
nii_files_ts = subfiles(test_dir, True, "test-volume", "nii", True)    # 获得test_dir中包含test-volume，且以nii结尾的数据

需要修改train_dir, test_dir, output_folder 三个路径，然后训练数据和训练标签在同一个文件夹下，只是命名不同。

运行改代码。就可以得到处理好的数据，存放在output_folder中。

3.2 修改路径

这一步需要给网络指定数据存放路径，代码为 nnUNet/nnunet/paths.py

# base = os.environ['nnUNet_raw_data_base'] if "nnUNet_raw_data_base" in os.environ.keys() else None
# preprocessing_output_dir = os.environ['nnUNet_preprocessed'] if "nnUNet_preprocessed" in os.environ.keys() else None
# network_training_output_dir_base = os.path.join(os.environ['RESULTS_FOLDER']) if "RESULTS_FOLDER" in os.environ.keys() else None

base = '/home2/*/nnunet_data'
preprocessing_output_dir = '/home2/*/nnunet_data/nnUNet_preprocessed'
network_training_output_dir_base = '/home2/*/nnunet_data/nnUNet_trained_models'

上方注释掉的原来的代码，需要在系统中设置。但是可以修改为下面的，直接指定具体地址。

if base is not None:
    nnUNet_raw_data = join(base, "nnUNet_raw_data")
    nnUNet_cropped_data = join(base, "nnUNet_cropped_data")
    maybe_mkdir_p(nnUNet_raw_data)
    maybe_mkdir_p(nnUNet_cropped_data)

从这段代码可以看出，base指向的刚才生成的 "nnUNet_raw_data"上级目录。因此可根据数据转换时候output_folder生成的路径进行修改。

3.3 预处理数据

该部分内容在 /nnUNet/nnunet/experiment_planning/nnUNet_plan_and_preprocess.py 中

可以到py文件所在的文件夹，运行命令：

python nnUNet_plan_and_preprocess.py -t XXX --verify_dataset_integrity
# -t: 选择任务名对应的ID
# --verify_dataset_integrity： “至少第一次在给定数据集上运行该命令时应该运行。这将对数据集执行一些检查，以确保它与nnU-Net兼容” 所以其实后来不加也可以

该命令行结束后即可得到“nnUNet_preprocessed/TaskXXX_MYTASK”文件夹。即处理后的数据。

3.4 训练

此处以2D UNet进行训练为例，该部分内容在 /nnUNet/nnunet/run/run_training.py 中

可以到py文件所在的文件夹，运行命令：

nnUNet_train 2d nnUNetTrainerV2 TaskXXX_MYTASK FOLD --npz
# 具体参数其实可以看run_training.py这个文件。
TaskXXX： 自己的任务ID和名字，例如 Task029_LiTS
FOLD:     0, 1, ..., 5 or 'all'
--npz：   是否保存中间softmax的输出
-c:       可以选择是否继续训练,同时可选择参数，使用哪个模型

3.5 测试

该部分内容在 /nnUNet/nnunet/inference/predict_simple.py 中

可以到py文件所在的文件夹，运行命令：

python predict_simple.py -i INPUT_FOLDER -o OUTPUT_FOLDER -t TASK_NAME_OR_ID -m CONFIGURATION --save_npz
# -chk: 选择模型名字
# 其他参数上面已经都讲了。也可以看predict_simple.py文件查看更多内容.

至此就可以在设置的 -o 文件夹中得到预测结果。

4 其他常见内容

次数介绍一下常见的需要修改的代码的内容。

4.1 修改显卡号

可以使用命令行修改，也可以在运行的py文件中指定

4.2 修改训练epoch

这需要先简单说一下文件调用。run_training.py中使用train_calss函数获得trainer, train_class又从default_configuration.py 中 get_default_configuration来，调用nnunet/training/network_training/nnUNetTrainerV2.py。继承nnUNetTrainer.py中函数nnUNetTrainer，继续继承network_trainer.py中NetworkTrainer。

因此可以查看etwork_trainer.py中436行 

while self.epoch < self.max_num_epochs:

得知修改： self.max_num_epochs。 在该文件夹直接修改不行，因为继承他的文件会覆盖改参数。因此在上层函数修改。

 

如有错误，多多指教；码字辛苦，感谢点赞！