输入含NaN时TensorFlow 2.0梯度为NaN的解决方法咨询

问题描述 #

我在构建简单回归模型时遇到了一个问题：输入数据中存在无法通过预处理替换的NaN值，使用TensorFlow 2.0的自动微分功能计算梯度时，结果返回了NaN。测试代码如下：

import tensorflow as tf
import numpy as np
x = np.float32(np.arange(100))
y = 2*x
x[0] = np.nan
a = tf.Variable([1e-6])
with tf.GradientTape() as t:
    t.watch(a)
    y_pred = a*x
    tensor = tf.reshape((y_pred - y), [-1])**2
    tensor_mask = tf.math.is_nan(tensor)
    tensor_without_nans = tf.where(tensor_mask, tf.zeros_like(tensor), tensor)
    RMSE = tf.reduce_sum(tensor_without_nans, axis=-1)
grads = t.gradient(RMSE,a)
print(grads)

问题原因 #

你当前的方法虽然在损失输出阶段用tf.where把NaN替换成了0，但TensorFlow的自动微分是基于整个计算图追踪梯度的。在你的计算路径中，y_pred = a*x已经产生了NaN值（因为x[0]是NaN），后续的平方、替换操作并没有消除计算图中NaN的传播路径，反向传播时NaN会一直传递到梯度结果中，导致最终得到NaN的梯度。

解决方案：提前过滤含NaN的样本 #

最直接有效的方法是在计算损失前就过滤掉含NaN的输入样本，让整个计算图只处理有效数值，这样梯度计算自然就不会出现NaN了。具体修改如下：

import tensorflow as tf
import numpy as np

x = np.float32(np.arange(100))
y = 2*x
x[0] = np.nan
a = tf.Variable([1e-6])

# 第一步：过滤出非NaN的有效样本
valid_mask = tf.math.is_finite(x)  # 标记所有非NaN、非无穷大的样本
x_valid = tf.boolean_mask(x, valid_mask)  # 提取有效输入
y_valid = tf.boolean_mask(y, valid_mask)  # 提取对应标签

with tf.GradientTape() as t:
    t.watch(a)
    y_pred = a * x_valid  # 仅对有效样本计算预测值
    # 直接计算有效样本的损失和（和你的RMSE逻辑一致，只是去掉了NaN样本）
    loss = tf.reduce_sum(tf.square(y_pred - y_valid))

grads = t.gradient(loss, a)
print(grads)

方案说明 #

• 我们通过tf.math.is_finite准确识别出所有有效样本（排除NaN和无穷大值），然后用tf.boolean_mask提取对应的输入和标签。
• 整个计算过程只涉及有效数值，计算图中不再有NaN相关的操作，反向传播时梯度就能正常计算并返回有效结果。
• 这个方法也符合回归任务的逻辑：含NaN的样本本身无法提供有效训练信号，直接排除是合理的选择。

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