训练简化版3通道U-Net时测试集CELoss为NaN的排查与解决

可能原因及对应解决办法： #

• 测试集数据存在异常
  训练集预处理正常但测试集可能混入了NaN/无穷值的体素数据，或者标签存在NaN、标签类别超出模型输出的类别范围（比如模型输出2类但标签出现3）。
  解决：遍历测试集所有样本，用torch.isnan(inputs).any()或np.isnan(labels).any()排查异常值；确认标签的取值范围严格匹配模型输出的类别数（CELoss要求标签为0到num_classes-1的整数）。

• 模型测试阶段输出极端值引发计算溢出
  训练时由于梯度裁剪、数据增强等约束，模型输出不会出现极端值，但单卷积层的简单模型泛化能力弱，测试时可能输出极大/极小的logits，经过softmax后出现0或1，计算log(0)时产生-inf，最终导致CELoss为NaN。
  解决：

  1. 对模型输出的logits做数值截断：logits = torch.clamp(logits, min=-100, max=100)，避免对数计算溢出；
  2. 给CELoss添加标签平滑：torch.nn.functional.cross_entropy(logits, labels, label_smoothing=1e-6)，软化极端的预测分布；
  3. 调整卷积层权重初始化：改用Xavier或He初始化（比如torch.nn.init.xavier_uniform_(conv.weight)），保证初始输出尺度合理。

• 训练/测试阶段模型模式不一致
  如果添加了批量归一化（BN）层，但测试时未切换模型到评估模式，BN层会继续使用当前测试batch的均值方差而非训练阶段统计的全局值，可能导致输出波动过大引发NaN。
  解决：测试前必须调用model.eval()，并用torch.no_grad()包裹前向传播代码，避免梯度计算干扰同时保证BN层正常工作。

• 测试集类别分布异常
  测试集中存在样本量极少的类别（甚至单个样本），模型对这类样本的预测完全偏离，导致交叉熵计算出现异常。
  解决：统计测试集标签分布，检查是否有类别样本数为0或极少；可以给CELoss添加类别权重（weight=class_weights，权重根据类别样本数反比设置），或者剔除测试集中的异常样本。

• 训练与测试数据预处理不一致
  测试集的归一化、缩放等预处理步骤和训练集不匹配，导致输入数据尺度异常，引发模型输出数值溢出。
  解决：对比训练和测试的预处理代码，确保均值、标准差、缩放比例等参数完全一致；检查是否存在测试集数据未做归一化的情况。

内容的提问来源于stack exchange，提问作者Pranjali Singh