Ubuntu 20.04 LTS下nnUNet训练医学影像数据集避坑指南(附完整代码)

📅 发布时间:2026/7/8 4:07:40 👁️ 浏览次数:
Ubuntu 20.04 LTS下nnUNet训练医学影像数据集避坑指南(附完整代码)
在Ubuntu 20.04 LTS上驯服nnUNet一份面向医学影像研究者的深度实战手册如果你正在医学影像分割的领域里摸索尤其是第一次接触那个传说中能“自适应”任何数据集的nnUNet框架那么你很可能正站在一个既令人兴奋又有点让人望而生畏的起点上。我完全理解这种感觉——面对一个功能强大但配置流程略显“固执”的工具从环境搭建到成功跑出第一个模型中间每一步都可能藏着意想不到的“坑”。这份指南就是为你准备的。它不是对官方文档的简单复述也不是一篇只讲“正确操作”的流水账。相反我会结合自己多次在Ubuntu 20.04 LTS上部署和调试nnUNet的实际经验把那些容易让人卡住数小时甚至数天的细节、错误信息背后的真实原因以及高效的排查思路毫无保留地分享出来。我们的目标很明确让你绕开那些常见的陷阱把精力真正集中在你的研究问题上。1. 环境奠基超越“pip install”的系统级准备很多人以为在Ubuntu上安装nnUNet一句pip install nnunet就万事大吉。这确实是开始但距离一个稳定、可用的训练环境还差着关键几步。特别是对于医学影像处理我们依赖的底层库对系统组件有特定要求。首先确保你的Ubuntu 20.04系统是最新的。打开终端执行sudo apt update sudo apt upgrade -y这能避免很多因系统包版本过旧导致的兼容性问题。接下来是CUDA和cuDNN。nnUNet的3D模型训练极度依赖GPU正确的CUDA环境是基石。假设你已安装NVIDIA驱动可以通过nvidia-smi查看驱动版本和最高支持的CUDA版本。对于Ubuntu 20.04CUDA 11.x系列是一个兼容性较好的选择。注意强烈建议使用官方.run文件安装CUDA而非apt仓库版本以避免与系统图形驱动产生冲突。安装时记得在安装选项中取消勾选附带的显卡驱动因为你已经安装好了。安装完成后将CUDA路径加入环境变量。编辑你的~/.bashrc文件nano ~/.bashrc在文件末尾添加请根据你的CUDA安装路径调整export PATH/usr/local/cuda-11.8/bin${PATH::${PATH}} export LD_LIBRARY_PATH/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH::${LD_LIBRARY_PATH}}保存后执行source ~/.bashrc并用nvcc --version验证。然后是nnUNet自身的环境变量这是核心配置决定了框架去哪里找你的数据和模型。同样在~/.bashrc中设置export nnUNet_raw_data_base/path/to/your/nnUNet_raw export nnUNet_preprocessed/path/to/your/nnUNet_preprocessed export RESULTS_FOLDER/path/to/your/nnUNet_trained_models这里的路径强烈建议设置为绝对路径并且不要包含中文或特殊字符。nnUNet_raw将存放你的原始数据集nnUNet_preprocessed存放预处理后的中间数据RESULTS_FOLDER则保存训练好的模型和日志。这三个目录需要提前手动创建好。最后创建一个独立的Python虚拟环境来安装nnUNet这是保持环境纯净的最佳实践conda create -n nnunet_env python3.8 -y conda activate nnunet_env pip install nnunet为什么是Python 3.8因为在我们的测试中它与PyTorch、nnUNet及其依赖的兼容性最为稳定。安装完成后可以运行nnUNet_print_available_pretrained_models来简单测试安装是否成功。2. 数据格式炼金术从“你的格式”到“nnUNet的格式”数据准备是使用nnUNet过程中最具挑战性的一环也是错误的高发区。框架对数据结构的严格要求是其实现自动化的前提但也成了新手的第一道门槛。关键在于理解其设计哲学而不仅仅是复制文件夹结构。核心规则一切围绕TaskXXX_任务名这个目录展开。这个目录必须位于你之前设置的nnUNet_raw_data_base目录下。XXX是三位数的任务ID它不仅是文件夹名的一部分更是后续所有命令中指代你这个数据集的“代号”。一个典型且无误的目录树应该如下所示/path/to/your/nnUNet_raw/ └── nnUNet_raw_data └── Task101_MyHeartMRI ├── imagesTr │ ├── patient001_0000.nii.gz │ ├── patient002_0000.nii.gz │ └── ... ├── labelsTr │ ├── patient001.nii.gz │ ├── patient002.nii.gz │ └── ... ├── imagesTs # 可选测试集 │ └── ... └── dataset.json这里有几个极易出错的魔鬼细节文件命名imagesTr中的图像文件必须是case_identifier_XXXX.nii.gz格式。case_identifier是病例标识符如patient001XXXX是四位数的模态标识符。例如对于一个多模态MRI数据集如T1和T2同一个病例的两个模态应命名为patient001_0000.nii.gzT1和patient001_0001.nii.gzT2。而对应的标签文件在labelsTr中必须且只能命名为patient001.nii.gz没有模态后缀。nii.gz文件内部确保你的.nii.gz文件不仅是这个扩展名其内部数据方向orientation、原点origin和间距spacing是合理的。有时从DICOM转换或其他软件导出的nifti文件其元信息affine matrix可能是错的这会导致nnUNet预处理时报出难以理解的形状错误。可以用NiBabel库简单检查import nibabel as nib img nib.load(‘patient001_0000.nii.gz’) print(img.affine) # 查看仿射矩阵 print(img.header[‘pixdim’]) # 查看体素间距dataset.json文件这是数据集的“身份证”。一个完整且正确的dataset.json示例及其关键字段解读如下字段名值示例说明与常见坑点“name”“MyHeartMRI”数据集名称可自定义。“description”“Cardiac MRI for ventricle segmentation.”描述有助于日后回顾。“tensorImageSize”“3D”必须是“3D”或“4D”。2D数据也选“3D”并将Z轴维度设为1。“modality”{“0”: “MRI”}指定每种模态的类型。键“0”对应_0000文件。如果是{“0”: “CT”}nnUNet会进行CT特定的预处理如裁剪强度值。务必与实际模态匹配。“labels”{“0”: “background”, “1”: “left_ventricle”, “2”: “myocardium”}定义标签中每个数值代表的类别。背景必须从0开始且键名是字符串格式的数字。“numTraining”100训练样本数需与training列表长度一致。“training”[{“image”: “./imagesTr/pat1_0000.nii.gz”, “label”: “./labelsTr/pat1.nii.gz”}, …]路径是相对于dataset.json文件本身的相对路径。这是最常见的错误来源之一绝对路径会导致失败。你可以用以下Python脚本作为模板来生成你的dataset.json这比手动编写更可靠import os import json from collections import OrderedDict # 设置路径 task_name “Task101_MyHeartMRI” raw_data_path f“/path/to/your/nnUNet_raw/nnUNet_raw_data/{task_name}” images_tr_path os.path.join(raw_data_path, “imagesTr”) labels_tr_path os.path.join(raw_data_path, “labelsTr”) # 获取所有训练图像ID去除模态后缀和扩展名 image_files [f for f in os.listdir(images_tr_path) if f.endswith(‘.nii.gz’)] # 假设是单模态提取唯一的case identifier case_ids list(set([f.split(‘_0000.nii.gz’)[0] for f in image_files])) case_ids.sort() # 构建JSON字典 dataset_json OrderedDict() dataset_json[‘name’] “MyHeartMRI” dataset_json[‘description’] “In-house cardiac MRI dataset.” dataset_json[‘tensorImageSize’] “3D” dataset_json[‘modality’] {“0”: “MRI”} dataset_json[‘labels’] {“0”: “background”, “1”: “target_organ”} dataset_json[‘numTraining’] len(case_ids) dataset_json[‘training’] [] for cid in case_ids: dataset_json[‘training’].append({ “image”: f“./imagesTr/{cid}_0000.nii.gz”, “label”: f“./labelsTr/{cid}.nii.gz” }) # 如果有测试集类似地添加 ‘numTest’ 和 ‘test’ 字段 # dataset_json[‘numTest’] … # dataset_json[‘test’] [“./imagesTs/xxx_0000.nii.gz”, …] # 写入文件 json_path os.path.join(raw_data_path, “dataset.json”) with open(json_path, ‘w’) as f: json.dump(dataset_json, f, indent4) print(f“dataset.json has been saved to: {json_path}”)3. 预处理与规划让nnUNet“认识”你的数据当数据格式无误后就可以运行nnUNet的自动化预处理流水线了。这个步骤会分析你的数据特性如体素间距、图像强度分布、形状大小并决定后续训练的最佳策略如是否使用patch-based训练、用什么分辨率等。命令很简单但背后的状态信息需要会看nnUNet_plan_and_preprocess -t 101 --verify_dataset_integrity-t 101指定你的任务ID就是Task101里的101。--verify_dataset_integrity首次运行时务必加上它会执行一系列检查。运行此命令时请紧盯终端输出这里列举几个常见的报错及解决方法错误Could not find ‘dataset.json’原因环境变量nnUNet_raw_data_base设置错误或者TaskXXX目录没有放在nnUNet_raw_data_base/nnUNet_raw_data/下面。解决用echo $nnUNet_raw_data_base检查环境变量。确保目录结构完全符合要求。错误Training case labelsTr/xxx.nii.gz does not have the same geometry as imagesTr/xxx_0000.nii.gz!原因图像和对应的标签文件体素网格尺寸、间距、方向不匹配。这是最常见也是最棘手的问题之一。解决使用nibabel或SimpleITK检查两个文件的shape和header信息。你可能需要重新用同一套参数来重采样或转换你的图像和标签数据确保它们对齐。警告Case xxx has more than 3 dimensions原因你的.nii.gz文件被加载后维度大于3可能是包含了时间维度或通道维度。解决医学影像分割通常处理3D体积数据。你需要将数据压缩或选择正确的维度。例如对于4D数据3D时间你可能需要提取某个时间点。如果预处理成功你会在nnUNet_preprocessed目录下找到Task101_MyHeartMRI文件夹里面包含了预处理后的数据和最重要的plans文件如nnUNetPlans.json。这个文件记录了nnUNet为你的数据集制定的“作战计划”包括目标分辨率、网络架构配置等。理解它有助于高级调优。4. 模型训练策略、监控与效率调优预处理通过后训练本身反而相对直接。但如何训练却有很多选择。基础训练命令# 训练3D全分辨率模型使用5折交叉验证的第0折 nnUNet_train 3d_fullres nnUNetTrainerV2 101 0这条命令会启动对任务ID 101第0折数据的训练。默认情况下它会训练1000个epoch这对于许多数据集可能过长。如何修改训练周期数直接修改源代码是最快的方式。找到你的nnUNet安装路径下的nnUNetTrainerV2.py文件。通常位于~/miniconda3/envs/nnunet_env/lib/python3.8/site-packages/nnunet/training/network_training/。找到__init__函数中的self.max_num_epochs 1000这一行将其改为你想要的数字例如self.max_num_epochs 300。提示更优雅的做法是创建自定义的训练器类来覆盖这个参数但对于快速实验直接修改是可行的。记得修改后重新启动训练。训练过程监控 训练开始后会在RESULTS_FOLDER/nnUNet/3d_fullres/Task101_MyHeartMRI/nnUNetTrainerV2__nnUNetPlans目录下或类似路径生成一个以时间戳命名的文件夹。里面你需要关注progress.png损失和验证集Dice系数曲线。这是判断模型是否收敛、是否过拟合的最直观依据。training_log.txt详细的训练日志包含每个epoch的损失值、学习率等。debug.json包含了所有训练配置的详细信息如果复现实验需要记录此文件。选择正确的训练配置 nnUNet提供了几种配置适用于不同特点的数据集2d2D U-Net。适用于切片间间距很大如2D超声序列或数据量极少的情况。训练快显存消耗小。3d_fullres3D U-Net在全分辨率图像上以patch形式训练。最常用的配置能捕捉3D上下文信息。3d_lowres和3d_cascade_fullres用于“级联U-Net”。当图像太大无法一次性放入GPU显存时先训练一个低分辨率模型进行粗分割再在裁剪出的感兴趣区域上用全分辨率模型进行细分割。这会显著增加训练和推理时间仅在必要时使用。一个实用的策略是先尝试3d_fullres。如果遇到GPU内存不足OOM错误再考虑使用3d_lowres或调整plans文件中的patch size。5. 推理与模型集成获得最终预测结果训练完成后五折交叉验证会生成五个模型你需要用这些模型对验证集或测试集进行预测并通常通过集成来提升鲁棒性。单模型预测# 使用训练好的3d_fullres模型对任务101的第0折模型进行推理 # -i 指定输入文件夹原始图像 # -o 指定输出文件夹 # -t 任务ID # -m 模型配置 # -f 指定使用哪一折的模型 (0,1,2,3,4) nnUNet_predict -i /path/to/test_images -o /path/to/predictions -t 101 -m 3d_fullres -f 0五折模型集成预测 这是nnUNet的默认最佳实践能有效提升性能nnUNet_predict -i /path/to/test_images -o /path/to/predictions -t 101 -m 3d_fullres不指定-f参数时命令会自动寻找五折模型并执行集成预测默认是平均软投票。预测后的结果处理 预测结果默认是软概率图每个类别的概率。如果你需要硬分割标签每个体素一个类别标签nnUNet在集成预测时会自动生成。你可以在输出文件夹中找到*.nii.gz每个测试病例的预测分割文件。postprocessing.json等如果启用了后处理如去除小连通域相关文件也会在这里。6. 高级排错与性能调优指南即使按照上述流程你可能还是会遇到一些独特的问题。这里分享几个“踩坑”后总结的经验。GPU内存不足OOM的解决方案减小Batch Size在自定义训练器或修改plans文件中可以尝试减小batch_size。但注意nnUNet的默认batch size通常已经很小如2。调整Patch Size这是最有效的方法。你需要修改预处理阶段生成的plans文件在预处理文件夹中。找到“patch_size”字段适当减小其值例如从[128, 128, 128]改为[96, 96, 96]。修改后需要重新运行nnUNet_plan_and_preprocess吗是的因为patch size是计划的一部分。更简单的方法是使用nnUNet_plan_and_preprocess时指定-plplan参数但更常见的做法是直接修改plans文件并重新预处理。使用3d_lowres配置如前所述级联模型的第一阶段就是在下采样图像上训练对显存要求更低。启用混合精度训练在自定义训练器中可以尝试启用AMP自动混合精度这能减少显存占用并可能加速训练。这需要修改训练器代码继承nnUNetTrainerV2并重写configure_optimizers和训练循环。训练损失震荡或不收敛检查数据首先确认标签是否正确。有时标签错误如全是背景会导致模型学不到东西。学习率nnUNet使用“poly”学习率衰减策略。如果损失曲线初期就剧烈震荡可能是初始学习率太大。你可以在自定义训练器中降低self.initial_lr。数据本身难度对于一些边界极其模糊或标注噪声很大的数据集Dice系数可能本身就不会很高。关注验证集损失是否在持续下降而不仅仅是最终指标。加速训练过程使用更快的存储将nnUNet_preprocessed和RESULTS_FOLDER放在SSD硬盘上能显著减少每个epoch的I/O等待时间。多GPU训练nnUNet原生支持DPDataParallel。在训练命令后加上-d参数指定GPU设备ID例如-d 0,1,2,3。但要注意多GPU训练可能会因为GPU间通信引入一些开销对于小batch size的情况加速比可能不理想。关闭验证集实时可视化在训练器中设置self.val_eval_criterion_alpha 0不推荐长期关闭因为无法监控过拟合。最后记住nnUNet社区是一个宝贵的资源。当你遇到一个模糊的错误信息时尝试在GitHub Issues或相关的学术论坛搜索错误关键词很大概率已经有人遇到过并解决了同样的问题。医学影像分割的道路从来不是一帆风顺每一个成功运行的模型背后都藏着几次与数据和代码的“搏斗”。希望这份指南能成为你手边的一盏灯照亮那些容易绊脚的石子路。