nnU-Net v2实战:从零开始配置环境与训练自定义医学影像数据集
1. 环境配置搭建nnU-Net v2的基础舞台第一次接触nnU-Net时我踩过的最大坑就是环境配置。当时为了赶项目进度直接用了现有的Python 3.8环境结果在安装时各种报错浪费了大半天时间。后来才发现nnU-Net v2对Python版本有严格要求必须≥3.9。关键步骤分解创建专用虚拟环境强烈推荐conda create -n nnunetv2 python3.9 conda activate nnunetv2这个操作相当于给你的项目单独开辟一个工作间避免与其他项目的依赖冲突。我习惯用conda管理环境用virtualenv也可以。PyTorch安装技巧 先运行nvidia-smi查看CUDA版本这个命令会显示类似CUDA Version:12.0的信息。记住这个数字然后去PyTorch官网选择匹配的安装命令。比如我的环境是CUDA 12.0就选择pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121验证安装 安装完成后别急着下一步先做个简单测试import torch print(torch.cuda.is_available()) # 应该返回True print(torch.__version__) # 确认版本≥1.12.0如果第一个输出是False说明GPU没启用可能是CUDA版本不匹配。2. nnU-Net框架安装不只是pip install那么简单很多人以为pip install nnunetv2就完事了其实这只是开始。完整安装需要三步走基础安装pip install nnunetv2源码克隆与开发模式安装关键git clone https://github.com/MIC-DKFZ/nnUNet.git cd nnUNet pip install -e .这个-e .操作让代码以可编辑模式安装后续可以直接修改源码调试。我在处理牙齿CT数据时就通过修改batchgenerators模块解决了内存溢出问题。可选组件 如果需要可视化网络结构pip install --upgrade githttps://github.com/FabianIsensee/hiddenlayer.git3. 数据集准备90%的坑都在这里处理医学影像数据集时文件结构比代码更重要。nnU-Net要求严格的目录规范nnUNet/ ├── DATASET/ │ ├── nnUNet_raw/ │ │ └── Dataset001_Teeth/ │ │ ├── imagesTr/ # 训练图像 │ │ ├── labelsTr/ # 训练标签 │ │ ├── imagesTs/ # 测试图像(可选) │ │ └── dataset.json # 关键配置文件 │ ├── nnUNet_preprocessed/ │ └── nnUNet_trained_models/dataset.json模板解析{ channel_names: {0: CBCT}, # 影像模态 labels: { # 标签定义 background: 0, Teeth: 1 }, numTraining: 12, # 训练样本数 file_ending: .nii.gz # 文件格式 }特别注意文件名必须遵循病例ID_000X.nii.gz格式比如case001_0000.nii.gz表示第一个病例的影像数据。4. 路径配置让nnU-Net找到你的数据遇到过最头疼的错误就是RuntimeError: Could not find dataset...都是路径配置不当导致的。推荐两种方法方法一环境变量配置永久生效修改~/.bashrc文件添加export nnUNet_raw/path/to/nnUNet_raw export nnUNet_preprocessed/path/to/nnUNet_preprocessed export nnUNet_results/path/to/nnUNet_trained_models然后执行source ~/.bashrc使配置生效。方法二运行时指定临时使用nnUNetv2_train DATASET_ID 3d_fullres 0 -pretrained_weights None --nnUNet_raw /path/to/raw --nnUNet_preprocessed /path/to/preprocessed --nnUNet_results /path/to/results5. 数据预处理nnU-Net的秘密武器执行这个命令开启自动化预处理nnUNetv2_plan_and_preprocess -d 1 --verify_dataset_integrity这里的1对应Dataset001中的ID号。预处理会完成重采样统一各向同性间距标准化CT值截断和归一化生成数据集指纹spacing、强度分布等统计信息在我的牙齿分割任务中预处理将原始CBCT数据体素间距0.3×0.3×0.3mm重采样到0.5×0.5×0.5mm显著降低了显存消耗。6. 模型训练五折交叉验证实战nnU-Net默认使用五折交叉验证有两种启动方式批量执行脚本推荐创建train.sh文件for fold in {0..4} do nnUNetv2_train 1 3d_fullres $fold done然后运行chmod x train.sh ./train.sh单折训练调试用nnUNetv2_train 1 3d_fullres 0 # 训练第1折训练中断后恢复nnUNetv2_train 1 3d_fullres 0 --c实用技巧使用nohup防止断连nohup nnUNetv2_train 1 3d_fullres 0 train.log 21 监控GPU使用watch -n 1 nvidia-smi7. 模型推理与应用训练完成后在nnUNet_results目录会生成类似Dataset001_Teeth/nnUNetTrainer__nnUNetPlans__3d_fullres的文件夹包含各折模型权重。单张影像预测nnUNetv2_predict -i input.nii.gz -o output.nii.gz -d 1 -f all -tr nnUNetTrainer其中-f all表示使用全部五折模型做ensemble。8. 实战经验牙齿分割特别注意事项在牙齿CBCT分割中我发现三个关键点金属伪影处理在dataset.json的channel_names中添加金属伪影标记通道小目标优化修改nnUNetPlans.json中的patch_size为[160,160,96]类别不平衡在训练命令添加--oversample_foreground_percent0.4对于牙根分割这种精细结构建议先用3d_lowres预训练再用3d_fullres微调。