PyTorch 如何计算二阶雅可比行列式？答案

【问题标题】：PyTorch how to compute second order Jacobian?PyTorch 如何计算二阶雅可比行列式？
【发布时间】：2020-11-23 23:42:09
【问题描述】：

我有一个神经网络正在计算一个向量 u。我想计算输入x（单个元素）的一阶和二阶雅可比。

有人知道如何在 PyTorch 中做到这一点吗？下面是我项目中的代码 sn-p：

import torch
import torch.nn as nn

class PINN(torch.nn.Module):
    
    def __init__(self, layers:list):
        super(PINN, self).__init__()
        self.linears = nn.ModuleList([])
        for i, dim in enumerate(layers[:-2]):
            self.linears.append(nn.Linear(dim, layers[i+1]))
            self.linears.append(nn.ReLU())
        self.linears.append(nn.Linear(layers[-2], layers[-1]))
        
    def forward(self, x):
        for layer in self.linears:
            x = layer(x)
        return x

然后我实例化我的网络：

n_in = 1
units = 50
q = 500

pinn = PINN([n_in, units, units, units, q+1])
pinn

PINN(
  (linears): ModuleList(
    (0): Linear(in_features=1, out_features=50, bias=True)
    (1): ReLU()
    (2): Linear(in_features=50, out_features=50, bias=True)
    (3): ReLU()
    (4): Linear(in_features=50, out_features=50, bias=True)
    (5): ReLU()
    (6): Linear(in_features=50, out_features=501, bias=True)
  )
)

然后我计算 FO 和 SO 雅可比

x = torch.randn(1, requires_grad=False)

u_x = torch.autograd.functional.jacobian(pinn, x, create_graph=True)
print("First Order Jacobian du/dx of shape {}, and features\n{}".format(u_x.shape, u_x)

u_xx = torch.autograd.functional.jacobian(lambda _: u_x, x)
print("Second Order Jacobian du_x/dx of shape {}, and features\n{}".format(u_xx.shape, u_xx)

First Order Jacobian du/dx of shape torch.Size([501, 1]), and features
tensor([[-0.0310],
        [ 0.0139],
        [-0.0081],
        [-0.0248],
        [-0.0033],
        [ 0.0013],
        [ 0.0040],
        [ 0.0273],
        ...
        [-0.0197]], grad_fn=<ViewBackward>)

Second Order Jacobian du/dx of shape torch.Size([501, 1, 1]), and features
tensor([[[0.]],

        [[0.]],

        [[0.]],

        [[0.]],

        ...

        [[0.]]])

如果不依赖于x，u_xx 不应该是None 向量吗？

提前致谢

【问题讨论】：

你检查过torch.autograd.functional.jacobian (pytorch.org/docs/stable/autograd.html)吗？
我有，而且我设法毫不费力地得到一阶雅可比，但我不知道如何计算二阶。我得到一个 0-valued ，考虑到我拥有的网络，这很奇怪。有没有办法可视化 PyTorch 动态生成的图表？
考虑到你的网络是由线性和 relu 组成的，它们几乎在任何地方都有恒定的梯度，所以输出 w.r.t 的二阶导数并不奇怪。输入处处为零。这对我来说似乎是正确的。
是的，我确实想通了，但还没有花时间回答自己！无论如何谢谢@jodag。我用nn.Tanh 替换了那些并得到了第二个订单。不过，我发现跟踪图表中填充的内容非常困难。计算雅可比两次是很棘手的。出于某种原因，torch.autograd.functional.jacobian 添加了无用的维度。这个我还没想好。
@aixyok 广义雅可比应该具有形状（输出形状 x 输入形状），因此第一个雅可比是 (501 x 1)，因为您的输入 x 大小为 1，输出 pinn 大小为 501。二阶Jacobain（又名 Hessian）将是 (501 x 1 x 1)，因为输出 u_x 的大小为 501 x 1，输入 x 的大小为 1。如果你愿意，可以在计算二阶雅可比之前 .flatten() u_x。

标签： python pytorch gradient

【解决方案1】：

所以正如@jodag 在他的评论中提到的，ReLU 为空或线性，它的梯度是恒定的（0 除外，这是一个罕见的事件），所以它的二阶导数为零。我将激活函数更改为Tanh，这最终允许我计算两次雅可比。

最终代码是

import torch
import torch.nn as nn

class PINN(torch.nn.Module):
    
    def __init__(self, layers:list):
        super(PINN, self).__init__()
        self.linears = nn.ModuleList([])
        for i, dim in enumerate(layers[:-2]):
            self.linears.append(nn.Linear(dim, layers[i+1]))
            self.linears.append(nn.Tanh())
        self.linears.append(nn.Linear(layers[-2], layers[-1]))
        
    def forward(self, x):
        for layer in self.linears:
            x = layer(x)
        return x
        
    def compute_u_x(self, x):
        self.u_x = torch.autograd.functional.jacobian(self, x, create_graph=True)
        self.u_x = torch.squeeze(self.u_x)
        return self.u_x
    
    def compute_u_xx(self, x):
        self.u_xx = torch.autograd.functional.jacobian(self.compute_u_x, x)
        self.u_xx = torch.squeeze(self.u_xx)
        return self.u_xx

然后在 PINN 的实例上调用 compute_u_xx(x) 并将 x.require_grad 设置为 True 让我到达那里。如何摆脱 torch.autograd.functional.jacobian 引入的无用维度仍有待理解...

【讨论】：

【解决方案2】：

二阶雅可比矩阵称为Hessian，可以使用 PyTorch 的内置函数轻松计算：

torch.autograd.functional.hessian(func, inputs)

【讨论】：

这仅适用于标量输出，不适用于向量