nnUNet V2实战：在AutoDL上从零构建医学图像分割工作流

1. 环境准备与AutoDL平台配置

在AutoDL云平台上部署nnUNet V2的第一步是搭建合适的环境。我推荐选择Ubuntu 20.04镜像，配合NVIDIA驱动版本≥470的GPU实例。具体操作中需要注意几个关键点：

1. GPU选择策略：对于3D医学图像分割，显存容量比核心数更重要。实测发现：

   • KiTS19肾脏肿瘤数据集（3D CT）在RTX 3090（24GB）上batch_size可设4
   • LiTS17肝脏数据集需要RTX 4090（24GB）才能流畅运行全分辨率模型

2. 存储配置技巧：

   bash复制# 将数据集放在/root/autodl-tmp（高速SSD）
   ln -s /root/autodl-tmp /dataset 
   # 模型输出指向/root/autodl-fs（持久化存储）
   mkdir -p /root/autodl-fs/nnUNet_results
   

3. Python环境快速搭建：

   bash复制conda create -n nnunet python=3.9
   conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 -c pytorch -c nvidia
   

提示：AutoDL的JupyterLab终端默认会保持会话状态，但建议用tmux管理长时间训练任务：

bash复制tmux new -s nnunet_train
# 按Ctrl+B D退出会话
tmux attach -t nnunet_train  # 恢复会话

2. nnUNet V2安装与验证

安装nnUNet V2时最容易踩坑的是路径配置问题。这里分享我的验证方法：

bash复制git clone https://github.com/MIC-DKFZ/nnUNet.git
cd nnUNet
pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

安装后必须设置三个关键环境变量（添加到~/.bashrc）：

bash复制export nnUNet_raw="/dataset/nnUNet_raw"
export nnUNet_preprocessed="/dataset/nnUNet_preprocessed" 
export nnUNet_results="/output/nnUNet_results"
source ~/.bashrc

验证安装成功的完整测试流程：

python复制import torch
from nnunetv2.run.run_training import run_training
print(torch.cuda.is_available())  # 应输出True

3. 医学数据集结构化处理

以KiTS19肾脏肿瘤数据集为例，需要严格遵循nnUNet的命名规范：

1. 目录结构：

   code复制nnUNet_raw/
   └── Dataset040_KiTS
       ├── imagesTr  # 训练图像
       ├── labelsTr  # 训练标签
       ├── imagesTs  # 测试图像
       └── dataset.json  # 元数据文件
   

2. dataset.json模板：

   json复制{
     "channel_names": {"0": "CT"},
     "labels": {
       "background": 0,
       "kidney": 1,
       "tumor": 2
     },
     "numTraining": 210,
     "file_ending": ".nii.gz"
   }
   

3. 自动转换脚本（适用于DICOM转NIfTI场景）：

   python复制import nibabel as nib
   from dicom2nifti import convert_directory
   convert_directory('/input/dicom', '/output/nifti')
   

4. 预处理与训练实战技巧

nnUNet V2的自动化预处理包含三个关键阶段：

1. 重采样优化：

   bash复制nnUNetv2_plan_and_preprocess -d 40 --verify_dataset_integrity
   

   • 自动计算最优体素间距（如KiTS19通常为3.22×1.62×1.62mm³）
   • 内存不足时可添加--lowres参数生成低分辨率版本

2. 训练参数调优：

   bash复制nnUNetv2_train 40 3d_fullres 0 -tr nnUNetTrainer_200epochs
   

   • 修改nnUNetTrainer.py中的self.num_epochs调整迭代次数
   • 使用--npz参数保存softmax输出用于模型集成

3. 多GPU训练技巧：

   bash复制CUDA_VISIBLE_DEVICES=0,1 nnUNetv2_train 40 3d_fullres 0 --use_compressed
   

5. 预测与结果后处理

nnUNet V2的预测流程支持多种实用场景：

1. 单模型预测：

   bash复制nnUNetv2_predict -i /input/imagesTs -o /output -d 40 -c 3d_fullres -f 0
   

2. 模型集成（5折交叉验证）：

   bash复制nnUNetv2_find_best_configuration 40 -c 2d 3d_fullres
   nnUNetv2_ensemble -i fold0 fold1 fold2 fold3 fold4 -o ensemble_results
   

3. 后处理优化：

   bash复制nnUNetv2_apply_postprocessing -i raw_predictions -o final_results \
   --pp_pkl_file /path/to/postprocessing.pkl
   

6. 性能评估与可视化

nnUNet提供两种评估方式：

1. 内置评估工具：

   bash复制nnUNetv2_evaluate -ref /labelsTs -pred /predictions -l 1 2
   

2. 自定义指标计算（Python实现）：

   python复制from nnunetv2.evaluation.evaluate_predictions import compute_metrics
   metrics = compute_metrics(reference_file, prediction_file, ['Dice', 'HD95'])
   

对于科研论文，推荐使用以下可视化组合：

• 3D Slicer 查看分割结果叠加
• ITK-SNAP 进行多平面对比
• Matplotlib 生成Dice分数分布直方图

7. 典型问题解决方案

内存不足错误：

python复制# 修改nnUNetPlans.json中的"patch_size"
"patch_size": [64, 160, 160]  # 原值[128, 160, 160]

训练中断恢复：

bash复制nnUNetv2_train 40 3d_fullres 0 --continue_training

预测结果异常检查清单：

1. 确认dataset.json中标签定义正确
2. 检查预处理是否完整执行
3. 验证GPU内存是否足够加载模型

在AutoDL平台上运行nnUNet V2时，建议定期使用nvidia-smi监控显存占用。对于大型数据集，可以分阶段执行预处理：

bash复制nnUNetv2_plan_and_preprocess -d 40 --no_pp
nnUNetv2_preprocess -d 40 -pl TrainerName -c 3d_fullres