torch.autograd.functional.jacobian

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

计算给定函数的雅可比矩阵。

参数:

• func (函数) – 一个接受张量输入并返回张量元组的 Python 函数。

• inputs (输入，张量元组或张量) – 函数 func 的输入。

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

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

• vectorize（布尔值，可选）- 此功能为实验性。如果需要更少实验性且性能更高的选项，请考虑使用 torch.func.jacrev() 或 torch.func.jacfwd() 。在计算雅可比矩阵时，通常我们每行雅可比矩阵调用 autograd.grad 一次。如果此标志为 True ，则仅使用 vmap 原型功能执行单个 autograd.grad 调用，其中使用 batched_grad=True 。尽管这可能在许多情况下提高性能，但由于此功能仍处于实验阶段，可能会出现性能悬崖。有关更多信息，请参阅 torch.autograd.grad() 的 batched_grad 参数。

• strategy（字符串，可选）- 设置为 "forward-mode" 或 "reverse-mode" 以确定是否使用前向或反向模式自动微分计算雅可比矩阵。目前， "forward-mode" 需要 vectorized=True 。默认为 "reverse-mode" 。如果 func 的输出比输入多，则 "forward-mode" 通常性能更高。否则，建议使用 "reverse-mode" 。

返回值:

如果只有一个输入和输出，这将是一个包含线性化输入和输出的雅可比的单一张量。如果其中两个之一是元组，则雅可比将是一个张量的元组。如果两个都是元组，则雅可比将是一个张量元组的元组，其中 Jacobian[i][j] 将包含 i 个输出和 j 个输入的雅可比，其大小为相应输出和相应输入大小的连接，并且具有与相应输入相同的 dtype 和 device。如果策略是 forward-mode ，则 dtype 将是输出的 dtype；否则，是输入的 dtype。

返回类型:

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

示例

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