如何防止TensorFlow训练x²预测模型时数据对≥18出现loss为inf

问题根源 #

输入xs（最大值30）和输出ys（最大值900）数值范围过大，且未做归一化就用SGD优化器训练，导致梯度计算值异常庞大，权重更新幅度过猛，很快触发数值溢出，最终loss变成inf。数据长度小于18时，ys最大值仅为289（17²），还没达到溢出阈值，所以能正常训练。

解决方案 #

1. 数据归一化（最关键）

将数据缩放到0-1区间，大幅降低梯度波动，避免数值溢出。预测时要同步对输入做归一化，再将结果反缩放回原范围：

import tensorflow as tf
import numpy as np
from tensorflow import keras

def quadratVonX(y_new):
    xs = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30], dtype=float)
    ys = np.array([0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196, 225, 256, 289, 324, 361, 400, 441, 484, 529, 576, 625, 676, 729, 784, 841, 900], dtype=float)
    
    # 归一化到0-1范围
    max_x = np.max(xs)
    max_y = np.max(ys)
    xs_scaled = xs / max_x
    ys_scaled = ys / max_y
    
    model = tf.keras.Sequential([keras.layers.Dense(units=1, input_shape=[1])])
    model.compile(optimizer='sgd', loss='mean_squared_error')
    model.fit(xs_scaled, ys_scaled, epochs=50, verbose=1)
    
    # 预测时同步归一化输入，再反缩放结果
    y_new_scaled = np.array(y_new) / max_x
    pred_scaled = model.predict(y_new_scaled, verbose=0)[0]
    return pred_scaled * max_y

prediction = quadratVonX([15])
print(prediction)

2. 更换优化器（可选，提升稳定性）

SGD对未归一化数据敏感度高，换成Adam优化器（自适应调整学习率），即使不做归一化，也能避免loss爆成inf：

model.compile(optimizer='adam', loss='mean_squared_error')

3. 优化模型结构（可选，提升拟合效果）

单神经元模型只能拟合线性关系，x²是非线性任务，加一层隐藏层能显著提升拟合精度（配合归一化效果更好）：

model = tf.keras.Sequential([
    keras.layers.Dense(units=8, input_shape=[1], activation='relu'),
    keras.layers.Dense(units=1)
])

效果验证 #

用归一化后的代码训练，loss会稳定下降，预测15的结果会接近真实值225，不会再出现inf的情况。

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