获取参数的梯度 w.r.t pytorch 中的损失项

Get grads of parameters w.r.t a loss term in pytorch

我在 Pytorch 训练中使用的复合损失函数定义为: 。 为了更新权重 alpha 和 beta,我需要计算三个值: 这是网络中所有权重的损失项 w.r.t 梯度的均值。

在pytorch中有没有高效的写法?

我的训练代码如下:

for epoc in range(1, nb_epochs+1):
  #init
  optimizer_fo.zero_grad()
  #get the current loss
  loss_total = mynet_fo.loss(tensor_xy_dirichlet,g_boundaries_d,tensor_xy_inside,tensor_f_inter,tensor_xy_neuman,g_boundaries_n)
  #compute gradients
  loss_total.backward(retain_graph=True)
  #optimize
  optimizer_fo.step()

我的 .loss() 函数直接 return 项的总和。我正在考虑进行第二次前向传递并独立地向后调用每个损失项,但这会非常昂贵。

分别获取每个损失分量的最佳方法是仅在模型外部定义损失(我现在假设损失是模型的一种方法,因为您将其称为方法)。

因此您应该将代码更改为类似于

ModelClass:
    def__init__(self):

    def forward(self):
       return output_of_the_model

# you can backpropagate inside the loss class directly 
LossClass(nn.Module):
    def__init__(self):

    def forward(self, model_output, target)
        loss_score_e = compute first component 
        loss_score_e.backward(retain_graph=True)
        # same for b and i components 
        return loss_score_e, loss_score_b, loss_score_i
 

那么训练脚本基本一样

loss = LossClass()
for epoc in range(1, nb_epochs+1):
      #init
      optimizer_fo.zero_grad()
      #get the current loss
      loss_e, loss_b, loss_i = loss(tensor_xy_dirichlet,g_boundaries_d,tensor_xy_inside,tensor_f_inter,tensor_xy_neuman,g_boundaries_n)
      #optimize
      optimizer_fo.step()

1- 使用 torch.autograd.grad

您只能通过在您的网络上多次反向传播来获得梯度的不同项。为了避免必须对您的输入执行多重推理,您可以使用 torch.autograd.grad 效用函数而不是执行传统的向后传递 backward。这意味着您不会污染来自不同项的梯度。

这是一个展示基本思想的最小示例:

>>> x = torch.rand(1, 10, requires_grad=True)
>>> lossA = x.pow(2).sum()
>>> lossB = x.mean()

然后对每个不合适的项执行一次反向传递。您必须保留除最后一次以外的所有调用的图表:

>>> gradA = torch.autograd.grad(lossA, x, retain_graph=True)
(tensor([[1.5810, 0.6684, 0.1467, 0.6618, 0.5067, 0.2368, 0.0971, 0.4533, 0.3511,
          1.9858]]),)

>>> gradB = torch.autograd.grad(lossB, x)
(tensor([[0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000,
          0.1000]]),)

此方法有一些限制,因为您接收参数的梯度作为 元组,这不是那么方便。


2-缓存backward

的结果

另一种解决方案是在每次连续 backward 调用后缓存梯度:

>>> lossA = x.pow(2).sum()
>>> lossB = x.mean()

>>> lossA.backward(retain_graph=True)

存储渐变并清除.grad属性(不要忘记这样做否则lossA的渐变会污染gradB。你将不得不适应这个处理多个张量参数时的一般情况:

>>> x.gradA = x.grad
>>> x.grad = None

下一个损失项的反向传播:

>>> lossB.backward()
>>> x.gradB = x.grad

然后你可以在本地与每个梯度项交互(分别在每个参数上):

>>> x.gradA, x.gradB
(tensor([[1.5810, 0.6684, 0.1467, 0.6618, 0.5067, 0.2368, 0.0971, 0.4533, 0.3511,
          1.9858]]),
 tensor([[0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000, 0.1000,
          0.1000]]))

后一种方法似乎更实用


这基本上归结为 torch.autograd.grad vs torch.autograd.backward, out-of-place vs in-place... 最终取决于您的需要。您可以阅读有关这两个函数的更多信息 .