如何计算二阶雅可比矩阵？

Question

我有一个计算向量的神经网络u。我想计算关于输入x的一阶和二阶雅可比矩阵，一个元素

有人知道在PyTorch怎么做吗？下面是我的项目中的代码片段：

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.]]])

如果u_xx不依赖于x，那么它不应该是None载体吗

提前谢谢

Answer 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

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

Answer 2

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

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