pytorch求二阶导数（loss函数中包含梯度信息） 2021-09-06

目前遇到了一个问题：如何在loss函数定义过程中加入梯度相关的内容。

如果直接进行backward与求grad的操作，会导致无法继续求导。例如，先算出结果相对于输入$\hat{h}$的导数，再算结果的模长相对于网络参数的导数：

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


class Classifier(nn.Module):
    def __init__(self, dim_in=5):
        super(DomainClassifier, self).__init__()
        self.fc = nn.Linear(dim_in, 1)

    def forward(self, x):
        x = self.fc(x)
        return x


if __name__ == '__main__':
    h_hat = torch.rand(13, 5).requires_grad_()
    d = Classifier()
    opt = torch.optim.Adam(d.parameters())
    opt.zero_grad()
    # 计算h_hat的梯度
    d(h_hat).sum().backward()
    grad = h_hat.grad
    L_grad = torch.norm(grad, p=2)
    print("L_grad", L_grad)
    # L_grad.backward() # 如果直接backward会报错
    # opt.step()

解决方案

1. 使用torch.autograd.grad函数：

   grad = torch.autograd.grad(d(h_hat).sum(), h_hat, create_graph=True)[0]
   print(grad)
   L_grad = torch.norm(grad, p=2)
   print("L_grad", L_grad)
   L_grad.backward()
   opt.step()
   

2. 第一个backward加上参数create_graph=True，并且后面算第二个导数之前加上zero_grad()：

   d(h_hat).sum().backward(create_graph=True)
   grad = h_hat.grad
   opt.zero_grad()
   L_grad = torch.norm(grad, p=2)
   L_grad.backward()
   opt.step()