TensorFlow入门示例扩展训练数据后输出NaN问题求助

1. 张量形状不匹配 #

原代码中xs和ys都是2D张量（形状[6,1]，代表6个样本，每个样本1个特征），但你的修改版用tf.tensor(d)生成的是1D张量，模型输入层定义为inputShape: [1]，要求输入必须是2D格式（每个样本是长度为1的向量），形状不匹配会导致训练计算出错，最终输出NaN。

2. 数据范围过大引发梯度爆炸 #

新数据的x范围是-50到9，数值远大于原示例的-1到4，而SGD优化器默认学习率（0.01）在处理大数值时，容易出现梯度更新幅度过大，导致权重变成NaN。

修改后的正确代码 #

async function run() {
  // Create a simple model.
  const model = tf.sequential();
  model.add(tf.layers.dense({units: 1, inputShape: [1]}));

  // 适当降低学习率，适配大数值范围
  model.compile({loss: 'meanSquaredError', optimizer: tf.train.sgd(0.001)});

  // Generate some synthetic data for training. (y = 2x - 1)
  const d = [];
  const r = [];
  for (let x = -50; x < 10; x++) {
    const y = x * 2 - 1;
    d.push(x);
    r.push(y);
  }
  // 保持2D张量形状，和原代码格式一致
  const xs = tf.tensor2d(d, [d.length, 1]);
  const ys = tf.tensor2d(r, [r.length, 1]);

  // Train the model using the data.
  await model.fit(xs, ys, {epochs: 250});

  // Use the model to do inference on a data point the model hasn't seen.
  // Should print approximately 39.
  document.getElementById('micro-out-div').innerText =
      model.predict(tf.tensor2d([20], [1, 1])).dataSync();
}

run();

额外优化建议 #

• 如果不想调整学习率，也可以对输入数据做归一化处理（比如把x缩放到-1到1之间），同样能避免梯度爆炸问题。
• 这类简单线性模型不需要盲目增加训练轮次，250轮已经足够，过多轮次反而可能在大学习率下导致震荡发散。

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