PyTorch梯度下降训练线性模型时持续出现NaN均方误差求助

在课程框架内尝试用PyTorch实现梯度下降拟合线性模型，编写的代码如下：

model = torch.nn.Linear(1,1)
myModel = model(X)
ds = torch.utils.data.TensorDataset(X, Y)
dl = torch.utils.data.DataLoader(ds)
optimiser = torch.optim.SGD(model.parameters(), lr=0.01)
loss = torch.nn.functional.mse_loss

for epoch in range(100):
    for (Xb, yb) in dl:
        yb_pred = model(Xb)
        c_loss = loss(yb_pred, yb)
        print(c_loss)
        optimiser.zero_grad()
        c_loss.backward()
        optimiser.step()

训练过程中持续输出NaN的MSE损失值，输出内容如下：

tensor(nan, grad_fn=)

问题分析与修正方案 #

以下是可能导致NaN损失的原因及对应解决方法：

• 数据未做归一化/标准化
  如果输入X或标签Y的数值范围过大（比如量级为1e3甚至1e6），模型参数更新时会引发梯度爆炸，直接导致损失计算出现NaN。解决方法是先对数据做标准化处理，将其缩放到均值为0、方差为1的范围：

  # 标准化X
  X = (X - X.mean()) / X.std()
  # 标准化Y（若Y范围过大也建议处理）
  Y = (Y - Y.mean()) / Y.std()
  

• 学习率设置不合理
  即使使用0.01的学习率，若数据未做归一化，仍可能导致参数更新幅度过大，触发梯度溢出。可以尝试降低学习率（比如调整为0.001），同时结合数据归一化使用，能有效避免NaN问题。

• 数据本身存在异常值
  检查输入X和标签Y中是否包含NaN或无穷大值，这类异常值会直接导致损失计算结果为NaN。可以用以下代码排查：

  print(torch.isnan(X).any())  # 检查X是否有NaN
  print(torch.isnan(Y).any())  # 检查Y是否有NaN
  print(torch.isinf(X).any())  # 检查X是否有Inf
  print(torch.isinf(Y).any())  # 检查Y是否有Inf
  

  若存在异常值，需要清理数据或替换为合理值。

• 冗余代码无意义
  代码中的myModel = model(X)属于冗余代码，这行仅在初始化模型后做了一次无意义的前向传播，既不参与训练也不影响结果，但建议删除以保持代码整洁。

修正后的示例代码 #

# 先处理数据标准化
X = (X - X.mean()) / X.std()
Y = (Y - Y.mean()) / Y.std()

model = torch.nn.Linear(1,1)
ds = torch.utils.data.TensorDataset(X, Y)
dl = torch.utils.data.DataLoader(ds)
optimiser = torch.optim.SGD(model.parameters(), lr=0.001)
loss = torch.nn.functional.mse_loss

for epoch in range(100):
    for (Xb, yb) in dl:
        yb_pred = model(Xb)
        c_loss = loss(yb_pred, yb)
        print(c_loss)
        optimiser.zero_grad()
        c_loss.backward()
        optimiser.step()

内容的提问来源于stack exchange，提问作者GaëtanLF