TensorFlow训练自编码器时损失返回NaN问题排查

AI 小北

2026-8-28

问题背景

正在训练输入为取值在[0,1]区间矩阵P的自编码器，所用加权交叉熵损失函数如下：
损失函数公式

损失函数的TensorFlow实现代码：

# Define loss and optimizer, minimize the squared error
with tf.device("/device:GPU:0"):
     L = -tf.reduce_sum(self.p*tf.math.log(self.p_pred+1e-10) + 0.55*(1 - self.p)*tf.math.log(1-self.p_pred+1e-10), axis = 1)
     self.loss = tf.reduce_mean(L)
self.optimizer = tf.compat.v1.train.AdamOptimizer(self.learning_rate).minimize(self.loss)

训练会话的运行逻辑：

_, l, y_pred, y = sess.run([model.optimizer, model.loss, model.y_pred, model.y], feed_dict=...)

当前训练配置：输出层激活函数为sigmoid，学习率设置为0.01，总训练轮次为20epoch。
故障现象：训练过程中损失会在随机epoch返回NaN。已尝试在损失出现NaN时保存对应的p和p_pred取值，在Google Colab中使用完全相同的损失函数计算这些保存值，结果不会出现NaN，无法定位问题原因。

排查思路

先区分NaN产生的阶段：你保存p、p_pred时拿到的是优化器步进完成后的数值，NaN大概率不是前向计算损失时产生的，而是出在反向传播梯度更新环节。sigmoid激活在输入绝对值过大时梯度本身趋近于0，0.01的学习率对于自编码器训练偏高，某次更新如果把网络层权重推到极大值，很容易触发梯度爆炸产生NaN，这时候你拿到的p_pred可能是更新后被内存改写/不同步的异常值，单独拿出来算前向损失自然复现不了问题。可以先把学习率降到1e-4跑2轮，验证NaN是否消失。
修复数值稳定性漏洞：你现在只在log的自变量位置加了1e-10的epsilon，这个操作只对前向计算的log有效，反向传播计算log的梯度公式为1/x，如果更新过程中p_pred取到0、1-p_pred取到0，梯度会直接变成inf，反向传导后直接把权重刷成NaN。不要只在log里加小量，在计算损失前先对预测值做硬裁剪：
```
self.p_pred = tf.clip_by_value(self.p_pred, 1e-7, 1 - 1e-7)
```
从源头截断超出[0,1]合法范围的预测值，同时把epsilon从1e-10调到1e-7，留足数值冗余。
排查GPU单精度计算的精度差问题：你的损失计算放在GPU上，CUDA核的单精度浮点数计算和CPU上的计算存在精度差异，某些临界值在GPU上计算时会先溢出为inf再传导成NaN，但你把值拷出到CPU环境复现计算时，精度差被抹平就不会触发异常。可以临时把损失计算放到CPU上跑几轮，如果不再出现NaN就可以确认是这个问题，除了上述的裁剪操作，还可以给Adam优化器加上epsilon=1e-7参数（默认值1e-8在单精度场景下稳定性不足）。
限制单步更新步长：Adam优化器通过累积梯度二阶矩自适应调整步长，如果某次训练batch存在特殊样本导致梯度出现尖峰，二阶矩估计未及时更新时，实际更新步长会远大于设置的0.01，直接把权重打飞。可以增加梯度全局裁剪逻辑，计算完梯度后先裁剪再应用到权重上，避免单步更新幅度过大：
```
opt = tf.compat.v1.train.AdamOptimizer(self.learning_rate)
grads_and_vars = opt.compute_gradients(self.loss)
grads, vars = zip(*grads_and_vars)
grads, _ = tf.clip_by_global_norm(grads, clip_norm=1.0)
self.optimizer = opt.apply_gradients(zip(grads, vars))
```
校验输入数据合法性：即使输入P理论取值在[0,1]，也要检查feed_dict喂入的批次是否偶尔混入NaN/inf值——比如数据预处理阶段多线程加载时的未初始化值、预处理逻辑的除0错误都可能产生异常值，这类异常值会直接把损失带成NaN，如果保存异常批次时刚好漏掉了含异常值的样本，就会出现复现失败的问题。可以在数据输入管道加断言，检查每个批次输入的最大值、最小值，以及是否存在NaN/inf值。

内容的提问来源于stack exchange，提问作者Hien Pham