torch 求Hessian矩阵

在 PyTorch 中可以使用 Autograd 模块来自动计算 Hessian 矩阵。具体来说，可以先计算目标函数关于参数的梯度，然后再对梯度求导，得到 Hessian 矩阵。

以下是一个简单的示例代码：

import torch

def hessian(y, x):
    """
    计算 y 对 x 的 Hessian 矩阵
    """
    # 计算 y 对 x 的一阶导数
    gradients = torch.autograd.grad(y, x, create_graph=True)
    # 初始化 Hessian 矩阵
    hessian = torch.zeros(x.numel(), x.numel())
    # 对梯度进行二阶求导
    for i, gradient in enumerate(gradients):
        # 将梯度展开成一列
        gradient_vector = gradient.view(-1)
        # 计算梯度向量对 x 的一阶导数
        jacobian = torch.zeros(gradient.numel(), x.numel())
        for j in range(gradient.numel()):
            jacobian[j] = torch.autograd.grad(gradient_vector[j], x, retain_graph=True)[0].view(-1)
        # 将 jacobian 矩阵的每一行放入 Hessian 矩阵
        hessian += jacobian.t() @ jacobian
    return hessian

其中，`y` 是目标函数，`x` 是待求 Hessian 矩阵的参数。这个函数返回的是一个二维张量，即 Hessian 矩阵。

需要注意的是，由于计算 Hessian 矩阵需要进行二阶求导，因此创建梯度张量时需要将 `create_graph` 参数设置为 `True`，以便保留计算图用于求导。此外，对于每个梯度向量，需要计算其对参数的一阶导数，这个过程也需要保留计算图，因此需要再次将 `retain_graph` 参数设置为 `True`。