nnUNet在2D心脏分割任务中损失值NaN问题的分析与解决

问题背景

在使用nnUNet框架进行心脏结构（左心室LV、心肌MYO和右心室RV）的2D分割任务时，开发者遇到了训练过程中损失值变为NaN（非数字）的问题。具体表现为训练损失和验证损失在几个epoch后突然变为NaN，同时Dice得分始终为0。值得注意的是，同样的数据在3D分割模式下表现正常。

现象分析

从训练日志可以看出：

1. 初始几个epoch损失值尚属正常范围（如0.36、0.12等）
2. 随着训练进行，损失值开始波动增大（如1.27、1.32等）
3. 约第6个epoch后，损失值突然变为NaN
4. Dice得分全程保持为0，说明模型未能学习到有效的分割特征

可能原因

1. 数据预处理问题

2D分割与3D分割对数据预处理的要求不同。可能存在的问题包括：

• 切片方向选择不当，导致解剖结构连续性被破坏
• 归一化方式不适合2D数据分布
• 2D切片中目标区域占比过小，导致类别极度不平衡

2. 训练策略差异

2D训练通常需要：

• 更小的初始学习率（相比3D）
• 不同的数据增强策略
• 特定的采样策略（如对包含目标的切片进行过采样）

3. 损失函数选择

默认的Dice+交叉熵损失在某些情况下可能导致数值不稳定，特别是：

• 当预测结果与真实标签完全不匹配时
• 存在极端类别不平衡时

解决方案

1. 使用专用训练器

建议尝试nnUNetTrainerDiceCELoss_noSmooth训练器，该版本：

• 移除了Dice损失中的平滑项
• 对极端情况有更好的数值稳定性

2. 调整训练参数

• 减小初始学习率（如从0.01降至0.001）
• 启用过采样（Oversampling）策略
• 减小patch size，特别是当目标区域较小时

3. 数据层面优化

• 检查2D切片方向是否合理（如选择短轴视图而非长轴视图）
• 验证标签在2D切片中的分布情况
• 考虑使用特定于心脏分割的数据增强策略

实施建议

1. 首先确保使用最新版本的nnUNet和相关依赖
2. 尝试最简单的配置（小学习率+noSmooth训练器）
3. 逐步添加复杂策略（如过采样）观察效果
4. 监控中间预测结果，确认模型是否在学习有意义的特征

总结

2D医学图像分割与3D分割在数据分布和训练动态上存在显著差异。针对心脏结构分割这一特定任务，需要特别注意类别平衡、损失函数选择和训练策略调整。通过合理配置训练器和优化参数，通常可以解决损失值NaN的问题，获得稳定的训练过程。