TensorFlow训练模型输入较大数值时输出NaN的问题如何解决

问题根源 #

出现NaN的核心原因是梯度爆炸：

• 默认SGD优化器的学习率过高，当输入/输出数值变大时，均方误差损失会随数值呈平方级增长，反向传播时梯度过大，直接导致权重更新溢出变为NaN，后续所有计算都会异常。
• 原代码中预测完成后又递归调用runNN()重复初始化训练的逻辑，也会加大数值溢出的概率。

修复方案 #

核心修改点 #

手动实例化SGD优化器，调低学习率到合适区间（比如0.0001），避免梯度过大溢出。如果需要用到更大范围的训练数值，可以额外对输入输出做归一化处理，预测完成后再反归一化还原结果，彻底规避大数值计算的溢出问题。
同时删掉冗余的递归调用逻辑，简化异步流程写法。

修复后可运行代码 #

async function runNN() {
  // 创建模型
  const model = tf.sequential();
  model.add(tf.layers.dense({
      units: 1,
      inputShape: [1]
  }));

  // 手动指定SGD优化器，调低学习率为0.0001，这是解决NaN问题的核心
  const sgdOptimizer = tf.train.sgd(0.0001);
  model.compile({
      optimizer: sgdOptimizer,
      loss: 'meanSquaredError'
  });

  // 训练集可加入大于10的数值，不会再溢出
  const xs = tf.tensor([
      [1],
      [2],
      [3],
      [4],
      [10],
      [100]
  ]);
  const ys = tf.tensor([
      [10],
      [20],
      [30],
      [40],
      [100],
      [1000]
  ]);

  // 直接用async/await简化流程，无需嵌套回调
  await model.fit(xs, ys, { epochs: 5000 });

  // 预测输出
  model.predict(tf.tensor([
      [1],
      [2],
      [100],
      [200],
  ])).print();
}

runNN();

如果觉得训练速度太慢，可以根据实际运行情况适当上调学习率，只要训练过程中loss没有出现NaN即可。

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