PyTorch 中自定义后向函数的损失 - 简单 MSE 示例中的爆炸损失

Posted 2023-03-12

【中文标题】PyTorch 中自定义后向函数的损失 - 简单 MSE 示例中的爆炸损失

【英文标题】：Loss with custom backward function in PyTorch - exploding loss in simple MSE example

【发布时间】：2021-05-02 22:45:56

【问题描述】：

在处理更复杂的事情之前，我知道我必须实现自己的backward 通行证，我想尝试一些简单而简单的事情。因此，我尝试使用 PyTorch 对均方误差损失进行线性回归。当我定义自己的 backward 方法时，出现了错误（请参阅下面的第三个实现选项），我怀疑这是因为我没有非常清楚地考虑将 PyTorch 作为渐变发送所需的内容。所以，我怀疑我需要的是一些关于 PyTorch 期望我在这里以什么形式提供的解释/澄清/建议。

我使用的是 PyTorch 1.7.0，所以一堆旧示例不再有效（使用用户定义的 autograd 函数的不同方式，如 documentation 中所述）。

第一种方法（标准 PyTorch MSE 损失函数）

让我们先用标准的方法来做，没有自定义损失函数：

import torch
import torch.nn as nn
import torch.nn.functional as F

# Let's generate some fake data
torch.manual_seed(42)
resid = torch.rand(100)    
inputs = torch.tensor([ [ xx ] for xx in range(100)] , dtype=torch.float32)
labels = torch.tensor([ (2 + 0.5*yy + resid[yy]) for yy in range(100)], dtype=torch.float32)

# Now we define a linear regression model
class linearRegression(torch.nn.Module):
    def __init__(self, inputSize, outputSize):
        super(linearRegression, self).__init__()
        self.bn = torch.nn.BatchNorm1d(num_features=1)
        self.linear = torch.nn.Linear(inputSize, outputSize)

    def forward(self, inx):
        x = self.bn(inx) # Adding BN to standardize input helps us use a higher learning rate
        x = self.linear(x)
        return x
    
model = linearRegression(1, 1)     

# Using the standard mse_loss of PyTorch
epochs = 25    
mseloss = F.mse_loss
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)

for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = mseloss(outputs.view(-1), labels)
    loss.backward()
    optimizer.step()
    scheduler.step()        
    print(f'epoch epoch, loss loss')    
    

这列火车很好，我损失了大约 0.0824，拟合图看起来很好。

第二种方式（自定义损失函数，但依赖PyTorch的自动梯度计算）

所以，现在我用自己的 MSE 损失实现替换了损失函数，但我仍然依赖 PyTorch autograd。我在这里唯一更改的是定义自定义损失函数，并在此基础上相应地定义损失，以及我如何将预测和真实标签交给损失函数的一个小细节。

#######################################################3
class MyMSELoss(nn.Module):
    
    def __init__(self):
        super(MyMSELoss, self).__init__()

    def forward(self, inputs, targets):        
        tmp = (inputs-targets)**2
        loss =  torch.mean(tmp)        
        return loss
#######################################################3

model = linearRegression(1, 1) 
    
mseloss = MyMSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)

for epoch in range(epochs):
    model.train()    
    outputs = model(inputs)
    loss = mseloss(outputs.view(-1), labels)
    loss.backward()    
    optimizer.step()
    optimizer.zero_grad()
    scheduler.step()
    
    print(f'epoch epoch, loss loss')

这给出了与使用标准 MSE 损失函数完全相同的结果。历元损失如下所示：

epoch 0, loss 884.2006225585938
epoch 1, loss 821.930908203125
epoch 2, loss 718.7732543945312
epoch 3, loss 538.1835327148438
epoch 4, loss 274.50909423828125
epoch 5, loss 55.115299224853516
epoch 6, loss 2.405021905899048
epoch 7, loss 0.47621214389801025
epoch 8, loss 0.1584305614233017
epoch 9, loss 0.09725229442119598
epoch 10, loss 0.0853077694773674
epoch 11, loss 0.08297089487314224
epoch 12, loss 0.08251354098320007
epoch 13, loss 0.08242412656545639
epoch 14, loss 0.08240655809640884
epoch 15, loss 0.08240310847759247
epoch 16, loss 0.08240246027708054
epoch 17, loss 0.08240233361721039
epoch 18, loss 0.08240240067243576
epoch 19, loss 0.08240223675966263
epoch 20, loss 0.08240225911140442
epoch 21, loss 0.08240220695734024
epoch 22, loss 0.08240220695734024
epoch 23, loss 0.08240220695734024
epoch 24, loss 0.08240220695734024

第三种方法（用我自己的后向方法自定义损失函数）

现在是最终版本，我在其中为 MSE 实现了我自己的渐变。为此，我在损失函数类中定义了自己的backward 方法，显然需要做mseloss = MyMSELoss.apply。

from torch.autograd import Function

#######################################################
class MyMSELoss(Function):
    
    @staticmethod
    def forward(ctx, y_pred, y):    
        ctx.save_for_backward(y_pred, y)
        return ( (y - y_pred)**2 ).mean()
    
    @staticmethod
    def backward(ctx, grad_output):
        y_pred, y = ctx.saved_tensors
        grad_input = torch.mean( -2.0 * (y - y_pred)).repeat(y_pred.shape[0])        
        # This fails, as does grad_input = -2.0 * (y-y_pred)
        # I've also messed around with the sign and that's not the sole problem, either.
        return grad_input, None
    
#######################################################
    
model = linearRegression(1, 1) 
mseloss = MyMSELoss.apply
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)

for epoch in range(epochs):
    model.train()
    outputs = model(inputs)
    loss = mseloss(outputs.view(-1), labels)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    scheduler.step()
    print(f'epoch epoch, loss loss')    

这是出了问题的地方，而不是训练损失减少，而是训练损失增加。现在看起来像这样：

epoch 0, loss 884.2006225585938
epoch 1, loss 3471.384033203125
epoch 2, loss 47768555520.0
epoch 3, loss 1.7422577779621402e+33
epoch 4, loss inf
epoch 5, loss nan
epoch 6, loss nan
epoch 7, loss nan
epoch 8, loss nan
epoch 9, loss nan
epoch 10, loss nan
epoch 11, loss nan
epoch 12, loss nan
epoch 13, loss nan
epoch 14, loss nan
epoch 15, loss nan
epoch 16, loss nan
epoch 17, loss nan
epoch 18, loss nan
epoch 19, loss nan
epoch 20, loss nan
epoch 21, loss nan
epoch 22, loss nan
epoch 23, loss nan
epoch 24, loss nan

【参考方案1】：

（2是常数可以忽略）

所以把你的后退函数改成这样：

@staticmethod
def backward(ctx, grad_output):
    y_pred, y = ctx.saved_tensors
    grad_input = 2 * (y_pred - y) / y_pred.shape[0]        
    return grad_input, None

【讨论】：

非常感谢，确实如此。为了向其他阅读本文的人说明，我们需要提供一个张量作为后向方法的输出，每个模型输出每个记录一个条目（即这里只是记录数，因此示例中的形状为 [100]，因为我们只对每条记录进行一次预测），其中包含相对于该记录的预测的梯度。 IE。在这种情况下，$m$ 记录中记录 $i$ 的预测梯度为 $2/m * (\haty_i - y_i)$。所以这最终是我们需要输出的张量的第 $i$ 个条目。