不用写代码神器！教你用4行命令轻松使用nnUNet训练自己的医学图像分割模型

Posted 2021-06-17 Tina姐

给定某个数据集，nnU-Net完全自动执行整个分割过程，包括数据预处理到模型配置、模型训练、后处理到集成的整个过程，而不需要人为干预。此外，训练好的模型还可以应用到测试集中进行推理。

博主强烈建议：做医学图像分割的任何人，都必须要会使用nnU-Net

理由2个：

• 首先用nnU-Net测试一下。看一下该任务的大致效果，心里有个大致预测。
• 第二，写论文的时候通常都要和nnU-Net对比，尤其是投英文期刊（SCI）。因为它太火了，审稿人都知道这个东西的，不要抱侥幸心理。要么战胜它，要么使用它😂。它的出现，真的是有好有坏，毕竟投论文更加难了~

开始之前，确保你已经学习了以下内容

论文解读- nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation（附实现教程）
nnU-Net 如何安装–傻瓜式安装教程
nnU-Net实战一使用预训练nnU-Net模型进行推理

文章目录

• 1 数据准备
• 2 一行命令数据预处理
• 3 一行命令开始训练
• 4 一行命令从断点开始训练
• 5 一行命令进行预测

1 数据准备

nnU-Net 需要结构化格式的数据集。并将原始数据存放在文件夹

nnUNet_raw_data_base/nnUNet_raw_data/TaskXXX_MYTASK

在nnU-Net实战一使用预训练nnU-Net模型进行推理 ，3.1 nnU-Net支持的数据格式这一小节已经讲了你需要把数据转换的格式。主要包含以下文件夹和文件。

• imagesTr: 训练图像。
• imagesTs（可选）: 测试图像
• labelTr: 训练集标签。
• dataset.json : 包含数据集的元数据， 如任务名字，模态，标签含义，训练集包含的图像地址👇。

这里讲一下，dataset.json这个文件怎么弄。一个合格的文件应该包含如下信息：

• name: 数据集名字
• dexcription: 对数据集的描述
• modality: 模态，0表示CT数据，1表示MR数据。nnU-Net会根据不同模态进行不同的预处理
• labels: label中，不同的数值代表的类别
• numTraining: 训练集数量
• numTest: 测试集数量
• training: 训练集的image 和 label 地址对
• test： 只包含测试集的image. 这里跟Training不一样
  
  现在我们知道长什么样了，那么怎么把我们的数据准备成这个样子呢。最简单的方法就是copy别人的代码，在前人的基础上修改一下。如： dataset_conversion

当然，代码能力强的可以自己写一个~

2 一行命令数据预处理

使用命令： nnUNet_plan_and_preprocess

nnUNet_plan_and_preprocess -t XXX --verify_dataset_integrity

作用： 它会自动读取 数据集的属性，例如图像大小、体素间距、强度信息等。并转换成模型需要的格式。

XXX是与您的任务名称相关联的整数标识符TaskXXX_MYTASK（如005， 就写5）。您可以一次传递多个任务 ID。

–verify_dataset_integrity 对数据集执行一些检查以确保它与 nnU-Net 兼容。如果此检查已通过一次，则可以在以后的运行中省略。如果您遵守数据集转换指南（见上文），那么这应该毫无问题地通过:-)

nnUNet_plan_and_preprocess完成后，结果保存在文件夹

nnUNet_preprocessed / TaskXXX_MYTASK。

主要包含以下文件：

3 一行命令开始训练

使用命令 nnUNet_train

该命令参数很多，可以使用nnUNet_train -h查看参数的含义

该命令的一般结构为：

nnUNet_train CONFIGURATION TRAINER_CLASS_NAME TASK_NAME_OR_ID FOLD (additional options)

• CONFIGURATION： 模型架构，三种Unet: 2D U-Net, 3D U-Net and a U-Net Cascade(U-Net级联)。
• TRAINER_CLASS_NAME： 使用的model trainer. 默认为nnUNetTrainerV2就可以
• TASK_NAME_OR_ID： 任务全名TaskXXX_MYTASK或者是ID号
• FOLD： 第几折交叉验证，可选 [0, 1, 2, 3, 4]，一共五折。

举例一： 2D U-Net 训练第三折

nnUNet_train 2d nnUNetTrainerV2 Task100_adrenal 3

或者

nnUNet_train 2d nnUNetTrainerV2 100 3

举例二： 3D full resolution U-Net 训练第二折

nnUNet_train 3d_fullres nnUNetTrainerV2 Task100_adrenal 2

举例三： 3D U-Net cascade 训练第二折

分为两步

step1 : 3D low resolution U-Net 训练

nnUNet_train 3d_lowres nnUNetTrainerV2 TaskXXX_MYTASK FOLD

step2: 3D full resolution U-Net 训练

nnUNet_train 3d_cascade_fullres nnUNetTrainerV2CascadeFullRes TaskXXX_MYTASK FOLD

tips： 3D full resolution U-Net 训练需要使用nnUNetTrainerV2CascadeFullRes

训练过程中，终端会输出以下信息：大致包括使用的模型，数据信息，损失函数，网络架构，每个epoch的输出。

我们将训练好的模型写入 RESULTS_FOLDER/nnUNet 文件夹。每次训练都会获得一个自动生成的输出文件夹名称：

nnUNet_preprocessed/CONFIGURATION/TaskXXX_MYTASKNAME/TRAINER_CLASS_NAME__PLANS_FILE_NAME/FOLD

文件树大致如下

每一折里面，又有图上红框所示内容。大致包括

• debug.json：包含用于训练此模型的蓝图和推断参数的摘要。不容易阅读，但对调试非常有用;-)
• model_best.model / model_best.model.pkl：训练期间确定的最佳模型的检查点文件。暂时没用。
• model_final_checkpoint.model / model_final_checkpoint.model.pkl：最终模型的检查点文件（训练结束后）。这是用于验证和推理的内容。
• progress.png：训练过程中的训练（蓝色）和验证（红色）损失图。还显示了评估指标的近似值（绿色）。这个近似值是前景类的平均 Dice 分数。
  
• validation_raw：在这个文件夹中是训练完成后预测的验证案例。

在训练期间，观察进度通常很有用。因此，我们随时可以查看 progress.png。它将在每个epoch后进行更新。

4 一行命令从断点开始训练

nnU-Net训练时间还是比较久的，毕竟默认跑1000epochs, 假如突然断电了怎么办，难道要重新来吗？

当然不是， nnU-Net 每 50 个 epoch 存储一个检查点。如果您需要继续之前的训练，只需在训练命令中添加一个==-c==即可。如：

nnUNet_train 2d nnUNetTrainerV2 100 3 -c

5 一行命令进行预测

训练完成后，我们当然要看看在测试集上的表现啦。

使用命令 nnUNet_predict

nnUNet_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -t TASK_NAME_OR_ID -m CONFIGURATION

INPUT_FOLDER: 测试数据地址
OUTPUT_FOLDER： 分割数据存放地址
CONFIGURATION： 使用的什么架构，2d or 3d_fullres or 3d_cascade_fullres等

如：

nnUNet_predict -i /home/.../nnunet_file/nnUNet_raw/nnUNet_raw_data/Task100_adrenal/imagesTs -o /home/.../nnunet_file/output -t 100 -m 2d -f 2

分割结果保存在OUTPUT_FOLDER，自行打开查看吧~

总结： 虽然写了这么多，其实也就4行命令就搞定了，每个步骤一行命令。nnUNet有一个不好的地方是，他的预测是没有Dice得分的，因为他不接收测试集的标签。

到这里，nnUNet的使用就分享结束了呀，欢迎评论区交流~

文章持续更新，可以关注微信公众号【医学图像人工智能实战营】获取最新动态，一个关注于医学图像处理领域前沿科技的公众号。坚持已实践为主，手把手带你做项目，打比赛，写论文。凡原创文章皆提供理论讲解，实验代码，实验数据。只有实践才能成长的更快，关注我们，一起学习进步~

我是Tina, 我们下篇博客见~

最后，求点赞，评论，收藏。或者一键三连