ResNet18训练首个epoch val_loss为NaN的问题排查与修复方案求助

优先尝试的快速修复方法 #

• 修改输出层与损失函数配置，提升数值稳定性：将ResNet18类最后一层的softmax激活去掉，改为 self.f3 = Dense(40)，同时修改损失函数配置为 tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)，TensorFlow会自动做softmax与交叉熵的融合计算，避免单独计算softmax时出现的数值溢出导致log(0)问题。
• 对齐训练集与验证集的预处理逻辑：确认验证集的图像输入x1_test、第二输入x2_test做了和训练集完全一致的预处理操作（包括归一化、标准化、数值缩放等），避免验证集输入数值范围过大导致激活值爆炸出现inf。
• 调整BatchNormalization参数：如果上述方法无效，可尝试将所有BN层的momentum参数调低至0.9，加快滑动平均统计值的更新速度，避免首个epoch训练结束后BN的滑动均值/方差还未收敛，导致验证阶段计算异常。也可临时注释所有BN层跑一轮验证，确认是否为BN层导致的问题。
• 补充校验验证集输入：除NaN外，额外检查x1_test、x2_test是否存在inf值，二次确认y_test的所有标签取值都在0~39范围内。

利用已保存Checkpoint定位NaN的方法 #

你可以通过手动逐batch跑验证前向的方式定位异常位置，步骤如下：

1. 加载模型权重：

model = Network_frame.ResNet18([2, 2, 2, 2])
model.load_weights('你的checkpoint文件路径')
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False) # 和你原来的配置保持一致

2. 开启数值异常检测：

tf.debugging.enable_check_numerics()

这个API会在计算过程中出现NaN/Inf时直接抛出异常，同时打印异常对应的算子和网络层位置。
3. 逐batch运行验证前向：

batch_size = 16
for batch_idx in range(len(x1_test) // batch_size):
    batch_x1 = x1_test[batch_idx*batch_size : (batch_idx+1)*batch_size]
    batch_x2 = x2_test[batch_idx*batch_size : (batch_idx+1)*batch_size]
    batch_y = y_test[batch_idx*batch_size : (batch_idx+1)*batch_size]
    # 用验证模式前向
    pred = model([batch_x1, batch_x2], training=False)
    loss = loss_fn(batch_y, pred)
    print(f"批次{batch_idx}损失：{loss.numpy()}")

如果某一批次出现loss为NaN，单独取出该批次的输入，逐层打印每层的输出值，即可定位到出现异常的网络层位置。

内容的提问来源于stack exchange，提问作者Paul Zhang