‘nnU-Net终极指南’
nnU-Net Ultimate Guide
了解最新的nnU-Net和如何将其应用于自己的数据集的一切
在我在剑桥大学进行深度学习和神经科学的研究实习期间，我经常使用nnU-Net，它是语义图像分割中极为强大的基准。
然而，我对这个模型以及如何训练它还有些困惑，并没有在互联网上找到太多帮助。现在我已经对它感到很舒适了，我创建了这个教程来帮助你，无论是在更好地了解这个模型背后的原理，还是如何在自己的数据集中使用它。
在本指南中，你将会：
- 对nnU-Net的关键贡献有一个简明的概述。
- 学习如何将nnU-Net应用于自己的数据集。
所有代码都可以在这个Google Collab笔记本上找到
这项工作花费了我大量的时间和精力。如果你觉得这个内容有价值，请考虑关注我以增加它的可见性，并支持更多类似教程的创作！
nnU-Net简史
nnU-Net被公认为图像分割中的最先进模型，在2D和3D图像处理方面都表现出色。它的性能非常强大，可以作为新的计算机视觉架构的强大基准。实际上，如果你要进入开发新的计算机视觉模型的世界，可以将nnU-Net视为你的“目标”。
这个强大的工具基于U-Net模型（你可以在这里找到我的一个教程：Cook your first U-Net），该模型于2015年首次亮相。名称“nnU-Net”代表“No New U-Net”，这是对其设计不引入革命性架构改变的一种致敬。相反，它采用现有的U-Net结构，并通过一系列巧妙的优化策略充分发挥其潜力。
与许多现代神经网络不同，nnU-Net不依赖于残差连接、稠密连接或注意机制。它的强大之处在于其细致的优化策略，包括重采样、归一化、损失函数的精选、优化器设置、数据增强、基于块的推断和模型集成等技术。这种整体方法使nnU-Net能够突破原始U-Net架构的可达到性极限。
探索nnU-Net中的多样化架构
尽管它可能看起来像一个单一的实体，但nnU-Net实际上是三种不同类型U-Net的总称：
- 2D U-Net：可以说是最著名的变体，直接处理2D图像。
- 3D U-Net：这是2D U-Net的扩展，通过应用3D卷积可以直接处理3D图像。
- U-Net级联：该模型生成低分辨率的分割结果，并进一步对其进行细化。
每种架构都有其独特的优势，但不可避免地也存在一定的局限性。
例如，使用2D U-Net进行3D图像分割可能看起来违反直觉，但实际上它仍然可以非常有效。这是通过将3D体积切片成2D平面来实现的。
虽然3D U-Net可能看起来更复杂，由于其更高的参数数量，它并不总是最高效的解决方案。特别是，3D U-Net在各向异性方面经常遇到困难，当空间分辨率沿不同轴线（例如，x轴上为1mm，z轴上为1.2mm）不同时会发生这种情况。
当处理大型图像尺寸时，U-Net级联变体变得特别有用。它使用一个初步模型来压缩图像，然后使用标准的3D U-Net输出低分辨率的分割。然后对生成的预测进行放大，从而得到精细、全面的输出。
通常，这种方法涉及在nnU-Net框架内训练所有三个模型变体。下一步可能是选择三个模型中表现最好的一个，或者使用集成技术。其中一种技术可能涉及整合2D和3D U-Net的预测。
然而，值得注意的是，这个过程可能非常耗时（也需要花费GPU资源）。如果你的限制只允许训练一个模型，不必担心。你可以选择只训练一个模型，因为集成模型只会带来非常微小的收益。
这个表格展示了与特定数据集相关的最佳模型变体：
网络拓扑的动态适应
由于图像尺寸存在显著差异（考虑肝脏图像的中位形状为482×512×512，而海马图像的形状为36×50×35），nnU-Net智能地调整输入块大小和每个轴的池化操作次数。这基本上意味着根据数据集自动调整卷积层的数量，方便有效地聚合空间信息。除了适应不同的图像几何结构外，该模型还考虑到可用内存等技术约束。
需要注意的是，该模型不直接在整个图像上执行分割，而是在仔细提取的具有重叠区域的补丁上执行。然后对这些补丁上的预测进行平均，得到最终的分割输出。
但是，较大的补丁意味着更多的内存使用，而批量大小也会消耗内存。所做的权衡是始终优先考虑补丁大小（模型的容量），而不是批量大小（只对优化有用）。
这是用于计算最佳补丁大小和批量大小的启发式算法：
以及对于不同数据集和输入尺寸的外观：
太好了！现在让我们快速回顾一下nnU-Net中使用的所有技术：
训练
所有模型都是从头开始训练，并使用训练集上的五折交叉验证进行评估，这意味着原始训练数据集被随机分为五个相等的部分或“折叠”。在这个交叉验证过程中，其中四个折叠用于训练模型，剩下的一个折叠用于评估或测试。这个过程被重复五次，每个折叠都被正好使用一次作为评估集。
对于损失函数，我们使用Dice和交叉熵损失的组合。这是图像分割中非常常见的损失函数。有关V-Net中Dice Loss的更多细节，可以参考U-Net的大兄弟。
数据增强技术
nnU-Net具有非常强大的数据增强流程。作者使用了随机旋转、随机缩放、随机弹性变形、gamma校正和镜像。
注意：您可以通过修改源代码添加自己的转换。
基于块的推理
正如我们所说，该模型不会直接在完整分辨率的图像上进行预测，而是在提取的块上进行预测，然后聚合预测结果。
具体如下：
注意：图片中心的块比侧边的块具有更高的权重，因为它们包含更多信息，模型对它们的性能更好。
成对模型集成
所以，如果您还记得，我们可以训练多达3个不同的模型，包括2D、3D和级联模型。但是在进行推理时，我们一次只能使用一个模型，对吗？
事实证明，不是这样。不同的模型具有不同的优势和劣势。因此，我们可以将多个模型的预测结果结合起来，以便如果一个模型非常自信，我们就优先考虑它的预测结果。
nnU-Net测试了可用模型中的2个模型的每种组合，并选择最佳的结果。
在实践中，有两种方法可以实现这一点：
硬投票：对于每个像素，我们查看两个模型输出的所有概率，并选择具有最高概率的类别。
软投票：对于每个像素，我们平均模型的概率，然后选择具有最大概率的类别。
实际实施
在开始之前，您可以在此处下载数据集，并按照Google Collab笔记本进行操作。
如果您对前面的部分一无所知，不用担心，这是实际部分，您只需跟着我做，就能获得最佳结果。
您需要有一个GPU来训练模型，否则它无法工作。您可以在本地或Google Collab上进行操作，不要忘记更改运行时> GPU
所以，首先，您需要准备好一个数据集，其中包含输入图像及其相应的分割。您可以按照我的教程下载这个准备好的3D脑部分割数据集，然后将其替换为您自己的数据集。
下载数据
首先，您应该下载数据并将其放置在data文件夹中，将两个文件夹分别命名为“input”和“ground_truth”，其中包含分割数据。
在接下来的教程中，我将使用MindBoggle数据集进行图像分割。您可以在此Google Drive上下载它：
我们获得了3D脑部的MRI扫描图像，我们希望对白质和灰质进行分割：
应该像这样：
设置主目录
如果在Google Colab上运行，请设置collab = True，否则设置collab = False
collab = Trueimport osimport shutil#librariesfrom collections import OrderedDictimport jsonimport numpy as np#visualization of the datasetimport matplotlib.pyplot as pltimport nibabel as nibif collab: from google.colab import drive drive.flush_and_unmount() drive.mount('/content/drive', force_remount=True) # 将"neurosciences-segmentation"更改为您的项目文件夹名称 root_dir = "/content/drive/MyDrive/neurosciences-segmentation"else: # 获取父目录的路径 root_dir = os.getcwd()input_dir = os.path.join(root_dir, 'data/input')segmentation_dir = os.path.join(root_dir, 'data/ground_truth')my_nnunet_dir = os.path.join(root_dir,'my_nnunet')print(my_nnunet_dir)现在我们要定义一个函数来为我们创建文件夹：
def make_if_dont_exist(folder_path,overwrite=False): """ 如果文件夹不存在，则创建文件夹 输入： folder_path：需要创建的文件夹的相对路径 over_write：（默认：False）如果为True，则覆盖现有文件夹 """ if os.path.exists(folder_path): if not overwrite: print(f'{folder_path}已存在。') else: print(f"{folder_path}被覆盖") shutil.rmtree(folder_path) os.makedirs(folder_path) else: os.makedirs(folder_path) print(f"{folder_path}已创建！")我们使用这个函数来创建我们的“my_nnunet”文件夹，这里将保存所有内容
os.chdir(root_dir)make_if_dont_exist('my_nnunet', overwrite=False)os.chdir('my_nnunet')print(f"当前工作目录：{os.getcwd()}")库安装
现在我们要安装所有的要求。首先让我们安装nnunet库。如果您在笔记本上运行，请在一个单元格中运行以下命令：
!pip install nnunet否则，您可以直接从终端安装nnunet：
pip install nnunet现在我们要克隆nnUnet git仓库和NVIDIA apex。这个仓库包含了训练脚本以及GPU加速器。
!git clone https://github.com/MIC-DKFZ/nnUNet.git!git clone https://github.com/NVIDIA/apex# repository dir is the path of the github folderrespository_dir = os.path.join(my_nnunet_dir,'nnUNet')os.chdir(respository_dir)!pip install -e!pip install --upgrade git+https://github.com/nanohanno/hiddenlayer.git@bugfix/get_trace_graph#egg=hiddenlayer创建文件夹
nnUnet对文件夹有非常特定的结构要求。
task_name = 'Task001' #更改此处以使用不同的任务名称#我们定义了所有必需的路径nnunet_dir = "nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data"task_folder_name = os.path.join(nnunet_dir,task_name) train_image_dir = os.path.join(task_folder_name,'imagesTr') #训练图像的路径train_label_dir = os.path.join(task_folder_name,'labelsTr') #训练标签的路径test_dir = os.path.join(task_folder_name,'imagesTs') #测试图像的路径main_dir = os.path.join(my_nnunet_dir,'nnUNet/nnunet') #主目录路径trained_model_dir = os.path.join(main_dir, 'nnUNet_trained_models') #训练模型的路径最初，nnU-Net是为了解决具有不同任务的十项全能挑战而设计的。如果您有不同的任务，只需为所有任务运行此代码块。
# 创建所有文件夹overwrite = False # 如果要覆盖文件夹，请将此参数设置为Truemake_if_dont_exist(task_folder_name, overwrite=overwrite)make_if_dont_exist(train_image_dir, overwrite=overwrite)make_if_dont_exist(train_label_dir, overwrite=overwrite)make_if_dont_exist(test_dir, overwrite=overwrite)make_if_dont_exist(trained_model_dir, overwrite=overwrite)现在您应该有如下结构：
设置环境变量
脚本需要知道您存放原始数据的位置、预处理数据的位置以及保存结果的位置。
os.environ['nnUNet_raw_data_base'] = os.path.join(main_dir, 'nnUNet_raw_data_base')os.environ['nnUNet_preprocessed'] = os.path.join(main_dir, 'preprocessed')os.environ['RESULTS_FOLDER'] = trained_model_dir将文件移动到正确的存储库中：
我们定义一个函数，将图像移动到nnunet文件夹中正确的存储库中：
def copy_and_rename(old_location, old_file_name, new_location, new_filename, delete_original=False): shutil.copy(os.path.join(old_location, old_file_name), new_location) os.rename(os.path.join(new_location, old_file_name), os.path.join(new_location, new_filename)) if delete_original: os.remove(os.path.join(old_location, old_file_name))现在让我们为输入图像和真值图像运行此函数：
list_of_all_files = os.listdir(segmentation_dir)list_of_all_files = [file_name for file_name in list_of_all_files if file_name.endswith('.nii.gz')]for file_name in list_of_all_files: copy_and_rename(input_dir, file_name, train_image_dir, file_name) copy_and_rename(segmentation_dir, file_name, train_label_dir, file_name)现在，我们需要将文件重命名为符合nnUnet格式的文件，例如subject.nii.gz将变为subject_0000.nii.gz
def check_modality(filename): """ 检查模态是否存在 如果模态不存在则返回False，否则返回True """ end = filename.find('.nii.gz') modality = filename[end-4:end] for mod in modality: if not(ord(mod)>=48 and ord(mod)<=57): #如果不是0到9的数字 return False return True def rename_for_single_modality(directory): for file in os.listdir(directory): if check_modality(file) == False: new_name = file[:file.find('.nii.gz')]+"_0000.nii.gz" os.rename(os.path.join(directory, file), os.path.join(directory, new_name)) print(f"重命名为 {new_name}") else: print(f"存在模态: {file}")rename_for_single_modality(train_image_dir)# rename_for_single_modality(test_dir)设置JSON文件
我们即将完成！
您主要需要修改两个内容：
- 模态（如果是CT或MRI，这将改变归一化）
- 标签：输入您自己的类别
overwrite_json_file = True #如果要覆盖Task_folder中的dataset.json文件，请将其设置为Truejson_file_exist = Falseif os.path.exists(os.path.join(task_folder_name, 'dataset.json')): print('dataset.json已存在！') json_file_exist = Trueif json_file_exist == False or overwrite_json_file: json_dict = OrderedDict() json_dict['name'] = task_name json_dict['description'] = "MindBoggle的T1扫描分割" json_dict['tensorImageSize'] = "3D" json_dict['reference'] = "参见挑战网站" json_dict['licence'] = "参见挑战网站" json_dict['release'] = "0.0" ######################## 修改此处 ######################## #您可以提及多个模态 json_dict['modality'] = { "0": "MRI" } #对数据集中的所有标签，都需要添加labels+1 json_dict['labels'] = { "0": "非脑部", "1": "皮层灰质", "2": "皮层白质", "3": "小脑灰质", "4": "小脑白质" } ############################################################# train_ids = os.listdir(train_label_dir) test_ids = os.listdir(test_dir) json_dict['numTraining'] = len(train_ids) json_dict['numTest'] = len(test_ids) #在dataset.json中，训练图像和标签没有模态 json_dict['training'] = [{'image': "./imagesTr/%s" % i, "label": "./labelsTr/%s" % i} for i in train_ids] #从测试图像名称中移除模态以保存在dataset.json中 json_dict['test'] = ["./imagesTs/%s" % (i[:i.find("_0000")]+'.nii.gz') for i in test_ids] with open(os.path.join(task_folder_name, "dataset.json"), 'w') as f: json.dump(json_dict, f, indent=4, sort_keys=True) if os.path.exists(os.path.join(task_folder_name, 'dataset.json')): if json_file_exist == False: print('dataset.json已创建！') else: print('dataset.json已覆盖！')为nnU-Net格式预处理数据
这将创建nnU-Net格式的数据集
# -t 1表示“Task001”，如果您有不同的任务，请更改它！nnUNet_plan_and_preprocess -t 1 --verify_dataset_integrity训练模型
我们现在准备好训练模型了！
要训练3D U-Net：
#训练3D全分辨率U-Net！nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --npz 要训练2D U-Net：
#训练2D U-Net！nnUNet_train 2d nnUNetTrainerV2 1 0 --npz要训练级联模型：
#训练3D U-Net级联！nnUNet_train 3d_lowres nnUNetTrainerV2CascadeFullRes 1 0 --npz!nnUNet_train 3d_fullres nnUNetTrainerV2CascadeFullRes 1 0 --npz注意：如果您暂停训练并希望恢复训练，请在末尾添加“-c”表示“继续”。
例如：
#训练3D全分辨率U-Net！nnUNet_train 3d_fullres nnUNetTrainerV2 1 0 --npz 推断
现在我们可以运行推断：
result_dir = os.path.join(task_folder_name, 'nnUNet_Prediction_Results')make_if_dont_exist(result_dir, overwrite=True)# -i是输入文件夹# -o是您想保存预测的位置# -t 1表示任务1，如果您有不同的任务编号，请更改它# 使用-m 2d，或-m 3d_fullres，或-m 3d_cascade_fullres！nnUNet_predict -i /content/drive/MyDrive/neurosciences-segmentation/my_nnunet/nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/Task001/imagesTs -o /content/drive/MyDrive/neurosciences-segmentation/my_nnunet/nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/Task001/nnUNet_Prediction_Results -t 1 -tr nnUNetTrainerV2 -m 2d -f 0 --num_threads_preprocessing 1预测结果可视化
首先让我们查看训练损失。这看起来非常健康，并且我们有一个Dice得分> 0.9（绿色曲线）。
对于这么少的工作和一个3D神经影像分割任务来说，这真的非常出色。
让我们看一个样本：
结果确实令人印象深刻！很明显，该模型已经有效地学会了如何高精度地分割脑部图像。虽然可能存在一些细微的缺陷，但重要的是要记住，图像分割领域正在快速发展，我们正朝着完美迈出重要的步伐。
将来，有可能进一步优化nnU-Net的性能，但那将是另一篇文章的内容。
如果您觉得本文有深入的见解并且有益，欢迎关注我以获取更深入的深度学习探索。您的支持帮助我继续制作有助于我们共同理解的内容。
无论您有什么反馈、想法要分享、想与我合作，还是只是想打个招呼，请填写下面的表格，让我们开始交流吧。
打个招呼 🌿
请毫不犹豫地给我点赞或关注我以获取更多内容！
参考资料
- Ronneberger, O., Fischer, P., & Brox, T. (2015). U-net: 用于生物医学图像分割的卷积网络. 在国际医学图像计算与计算机辅助干预会议上 (pp. 234–241). Springer, Cham.
- Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: 一种基于深度学习的自配置方法用于生物医学图像分割. Nature Methods, 18(2), 203–211.
- Ioffe, S., & Szegedy, C. (2015). 批归一化: 通过减少内部协变量偏移来加速深度网络训练. arXiv预印本 arXiv:1502.03167.
- Ulyanov, D., Vedaldi, A., & Lempitsky, V. (2016). 实例归一化: 快速风格化的缺失要素. arXiv预印本 arXiv:1607.08022.
- MindBoggle 数据集
