torch.autograd.functional.hessian

torch.autograd.functional.hessian(func, inputs, create_graph=False, strict=False, vectorize=False, outer_jacobian_strategy='reverse-mode')

计算给定标量函数的 Hessian 矩阵。

参数

• func (函数) – 一个接受 Tensor 输入并返回单个元素的 Tensor 的 Python 函数。

• inputs (tuple of Tensors 或 Tensor) – 函数 func 的输入。

• create_graph (bool, 可选) – 如果为 True，将以可微分的方式计算 Hessian 矩阵。请注意，当 strict 为 False 时，结果不能需要梯度或与输入断开连接。默认为 False。

• strict (bool, 可选) – 如果为 True，当我们检测到某个输入使得所有输出都与其无关时，将引发错误。如果为 False，我们将返回一个零 Tensor 作为这些输入的 Hessian 矩阵，这是预期的数学值。默认为 False。

• vectorize (bool, 可选) – 此功能尚处于实验阶段。如果您正在寻找一个更成熟且性能更好的替代方案，请考虑使用 torch.func.hessian()。计算 Hessian 矩阵时，通常我们需要对 Hessian 矩阵的每一行调用一次 autograd.grad。如果此标志为 True，我们将使用 vmap 原型功能作为后端来向量化对 autograd.grad 的调用，从而只需调用一次，而非每行调用一次。这在许多使用场景下应能带来性能提升，但是，由于此功能尚不完善，可能会出现性能瓶颈。请使用 torch._C._debug_only_display_vmap_fallback_warnings(True) 来显示任何性能警告，如果您在使用中遇到警告，请向我们提交问题。默认为 False。

• outer_jacobian_strategy (str, 可选) – Hessian 矩阵通过计算 Jacobian 矩阵的 Jacobian 矩阵来获得。内部 Jacobian 矩阵总是使用反向模式 AD 计算。将 strategy 设置为 "forward-mode" 或 "reverse-mode" 决定外部 Jacobian 矩阵是使用前向模式 AD 还是反向模式 AD 计算。目前，使用 "forward-mode" 计算外部 Jacobian 矩阵需要设置 vectorized=True。默认为 "reverse-mode"。

返回

如果输入为单个，则返回一个包含该输入 Hessian 矩阵的单个 Tensor。如果输入为一个 tuple，则返回一个 tuple 的 tuple，其中 Hessian[i][j] 包含第 i 个输入和第 j 个输入的 Hessian 矩阵，其大小等于第 i 个输入大小与第 j 个输入大小之和。Hessian[i][j] 的 dtype 和设备将与对应的第 i 个输入相同。

返回类型

Hessian (Tensor 或由 Tensors 的 tuple 组成的 tuple)

示例

>>> def pow_reducer(x):
...     return x.pow(3).sum()
>>> inputs = torch.rand(2, 2)
>>> hessian(pow_reducer, inputs)
tensor([[[[5.2265, 0.0000],
          [0.0000, 0.0000]],
         [[0.0000, 4.8221],
          [0.0000, 0.0000]]],
        [[[0.0000, 0.0000],
          [1.9456, 0.0000]],
         [[0.0000, 0.0000],
          [0.0000, 3.2550]]]])

>>> hessian(pow_reducer, inputs, create_graph=True)
tensor([[[[5.2265, 0.0000],
          [0.0000, 0.0000]],
         [[0.0000, 4.8221],
          [0.0000, 0.0000]]],
        [[[0.0000, 0.0000],
          [1.9456, 0.0000]],
         [[0.0000, 0.0000],
          [0.0000, 3.2550]]]], grad_fn=<ViewBackward>)

>>> def pow_adder_reducer(x, y):
...     return (2 * x.pow(2) + 3 * y.pow(2)).sum()
>>> inputs = (torch.rand(2), torch.rand(2))
>>> hessian(pow_adder_reducer, inputs)
((tensor([[4., 0.],
          [0., 4.]]),
  tensor([[0., 0.],
          [0., 0.]])),
 (tensor([[0., 0.],
          [0., 0.]]),
  tensor([[6., 0.],
          [0., 6.]])))