torch.autograd.functional.jacobian

torch.autograd.functional.jacobian(func, inputs, create_graph=False, strict=False, vectorize=False, strategy='reverse-mode')[源代码]

计算给定函数的雅可比行列式。

参数

func（函数）– 一个接受张量输入并返回张量元组或单个张量的 Python 函数。
inputs (元组 of Tensors 或 Tensor) – 传递给函数 func 的输入。
create_graph (bool, optional) – 如果设置为 True，则以可微分的方式计算雅可比矩阵。需要注意的是，当 strict 为 False 时，结果不能需要梯度或与输入断开连接。默认值为 False。
strict (bool, optional) – 如果设置为 True，当检测到存在某个输入使得所有输出都与其无关时会引发错误。如果设置为 False，对于所述的输入返回一个全零张量作为雅可比矩阵，这是数学上的预期值。默认为 False。
vectorize (bool, 可选) – 此功能尚处于实验阶段。如果你需要更稳定且性能更好的方案，请考虑使用torch.func.jacrev() 或者 torch.func.jacfwd()。在计算雅可比矩阵时，通常我们每次调用一次autograd.grad 来生成雅可比矩阵的一行。如果此标志为 True，我们将只进行一次带有 batched_grad=True 的 autograd.grad 调用，这使用了 vmap 原型功能。尽管在许多情况下这可能会带来性能提升，但由于此功能仍处于实验阶段，可能存在性能瓶颈。有关更多信息，请参阅torch.autograd.grad() 的 batched_grad 参数。
strategy (str, optional) – 设置为"forward-mode"或"reverse-mode"以确定雅可比矩阵是使用前向模式还是反向模式自动微分来计算。目前，选择"forward-mode"需要设置vectorized=True。默认值为"reverse-mode"。如果函数func的输出多于输入，则通常推荐使用"forward-mode"以获得更好的性能。否则，建议使用"reverse-mode"。

返回值

如果有单一的输入和输出，雅可比矩阵将会是一个包含线性化输入和输出结果的单个张量。如果其中一方是元组，则雅可比矩阵会变成一个张量组成的元组。若两者都是元组，那么雅可比矩阵将是一个由张量组成的元组集合，其中Jacobian[i][j]表示第i个输出和第j个输入的雅可比矩阵，并且其大小等于相应输出与输入尺寸的连接。数据类型和设备会与相应的输入一致。如果策略是forward-mode，则数据类型将依据输出确定；否则，依据输入。

返回类型

雅可比矩阵（张量或张量的嵌套元组）

示例

>>> def exp_reducer(x):
...     return x.exp().sum(dim=1)
>>> inputs = torch.rand(2, 2)
>>> jacobian(exp_reducer, inputs)
tensor([[[1.4917, 2.4352],
         [0.0000, 0.0000]],
        [[0.0000, 0.0000],
         [2.4369, 2.3799]]])

>>> jacobian(exp_reducer, inputs, create_graph=True)
tensor([[[1.4917, 2.4352],
         [0.0000, 0.0000]],
        [[0.0000, 0.0000],
         [2.4369, 2.3799]]], grad_fn=<ViewBackward>)

>>> def exp_adder(x, y):
...     return 2 * x.exp() + 3 * y
>>> inputs = (torch.rand(2), torch.rand(2))
>>> jacobian(exp_adder, inputs)
(tensor([[2.8052, 0.0000],
        [0.0000, 3.3963]]),
 tensor([[3., 0.],
         [0., 3.]]))