使用PyTorch实现一元线性回归时出现NaN输出问题

当训练数据包含较大整数（比如把训练样本从(20,40)换成(50,100)）时，模型训练后预测输出变为NaN，而使用较小数值时预测正常。

实现代码 #

import torch
from torch.autograd import Variable
import numpy as np

class linearRegression(torch.nn.Module):            
    
    # instantiating Linear Regression class
    def __init__(self):
        super(linearRegression, self).__init__() 
        self.lr_pytorch = torch.nn.Linear(1,1)                 
    def forward(self,input_to_predict):                        
        predictions = self.lr_pytorch(input_to_predict)
        return predictions
            
    def fit(self,inputs,outputs,epochs = 10, learning_rate = 0.01):  
        
        inputs = Variable(torch.from_numpy(np.array(inputs,dtype=np.float32).reshape(-1,1)))        
        labels = Variable(torch.from_numpy(np.array(outputs,dtype=np.float32).reshape(-1,1))) 
        
        stochastic_gradient_descent = torch.optim.SGD(self.lr_pytorch.parameters(),learning_rate)  
        mse_loss_function = torch.nn.MSELoss()                                                     
        
        for epoch in range(epochs):                                                                 
            stochastic_gradient_descent.zero_grad()
            predictions = self.lr_pytorch(inputs)
            
            loss = mse_loss_function(predictions,labels)
            loss.backward()
            stochastic_gradient_descent.step()

测试案例对比 #

正常工作的训练数据 #

a = [1,2,3,4,5,6,7,8,9,10,20]
b = [2,4,6,8,10,12,14,16,18,20,40]
c = Variable(torch.from_numpy(np.array([150],dtype=np.float32)))

d = linearRegression()
d.fit(a,b,500)

prediction = d.forward(c) 

print(prediction)

输出：tensor([299.9434], grad_fn=<AddBackward0>)

出现NaN的训练数据 #

a = [1,2,3,4,5,6,7,8,9,10,50]
b = [2,4,6,8,10,12,14,16,18,20,100]
c = Variable(torch.from_numpy(np.array([150],dtype=np.float32)))

d = linearRegression()
d.fit(a,b,500)

prediction = d.forward(c) 

print(prediction)

输出：tensor([nan], grad_fn=<AddBackward0>)

解决方法 #

1. 对输入数据做归一化/标准化 #

大数值输入会导致MSE损失值过大，进而引发梯度爆炸，参数更新后直接变成NaN。将数据缩放到较小范围（比如均值为0、方差为1的标准分布），能从根源上避免梯度异常。

修改fit方法加入归一化逻辑，同时调整forward方法支持反归一化：

class linearRegression(torch.nn.Module):            
    def __init__(self):
        super(linearRegression, self).__init__() 
        self.lr_pytorch = torch.nn.Linear(1,1)                 
        # 初始化归一化参数
        self.input_mean = 0.0
        self.input_std = 1.0
        self.output_mean = 0.0
        self.output_std = 1.0
        
    def forward(self,input_to_predict):                        
        # 对输入做归一化
        input_normalized = (input_to_predict - self.input_mean) / self.input_std
        predictions_normalized = self.lr_pytorch(input_normalized)
        # 反归一化得到真实预测值
        predictions = predictions_normalized * self.output_std + self.output_mean
        return predictions
            
    def fit(self,inputs,outputs,epochs = 10, learning_rate = 0.01):  
        inputs_np = np.array(inputs,dtype=np.float32).reshape(-1,1)
        outputs_np = np.array(outputs,dtype=np.float32).reshape(-1,1)
        
        # 计算归一化参数
        self.input_mean = np.mean(inputs_np)
        self.input_std = np.std(inputs_np)
        self.output_mean = np.mean(outputs_np)
        self.output_std = np.std(outputs_np)
        
        # 归一化处理
        inputs_normalized = (inputs_np - self.input_mean) / self.input_std
        outputs_normalized = (outputs_np - self.output_mean) / self.output_std
        
        inputs = Variable(torch.from_numpy(inputs_normalized))        
        labels = Variable(torch.from_numpy(outputs_normalized)) 
        
        stochastic_gradient_descent = torch.optim.SGD(self.lr_pytorch.parameters(),learning_rate)  
        mse_loss_function = torch.nn.MSELoss()                                                     
        
        for epoch in range(epochs):                                                                 
            stochastic_gradient_descent.zero_grad()
            predictions = self.lr_pytorch(inputs)
            
            loss = mse_loss_function(predictions,labels)
            loss.backward()
            stochastic_gradient_descent.step()

2. 降低学习率 #

原学习率0.01在大数值场景下可能过高，导致参数更新幅度过大。可以尝试将学习率调低到0.001甚至更小：

d.fit(a,b,500, learning_rate=0.001)

3. 添加梯度裁剪 #

在反向传播后限制梯度的最大范数，防止梯度爆炸：
修改fit方法中的循环部分：

for epoch in range(epochs):                                                                 
    stochastic_gradient_descent.zero_grad()
    predictions = self.lr_pytorch(inputs)
    
    loss = mse_loss_function(predictions,labels)
    loss.backward()
    # 梯度裁剪，限制最大范数为1.0
    torch.nn.utils.clip_grad_norm_(self.lr_pytorch.parameters(), max_norm=1.0)
    stochastic_gradient_descent.step()

优先推荐数据归一化方案，它不仅解决NaN问题，还能提升模型训练的稳定性和收敛速度。

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