torch.autograd.functional.jacobian

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

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

参数

• func (function) – 一个 Python 函数，接受 Tensor 输入并返回一个 Tensor 元组或一个 Tensor。

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

• create_graph (bool, optional) – 如果设置为 True，则雅可比矩阵的计算将是可微的。请注意，当 strict 为 False 时，结果不能要求梯度或与输入断开连接。默认为 False。

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

• vectorize (bool, optional) – 此功能是实验性的。如果您正在寻找一些不太实验性且性能更好的功能，请考虑改用 torch.func.jacrev() 或 torch.func.jacfwd()。在计算雅可比矩阵时，通常我们会为雅可比矩阵的每一行调用一次 autograd.grad。如果此标志为 True，我们将只进行一次 autograd.grad 调用，并设置 batched_grad=True，该调用使用 vmap 原型功能。虽然这在许多情况下应能提高性能，但由于此功能仍处于实验阶段，因此可能存在性能瓶颈。有关更多信息，请参阅 torch.autograd.grad() 的 batched_grad 参数。

• strategy (str, optional) – 设置为 "forward-mode" 或 "reverse-mode" 以确定雅可比矩阵将使用前向模式或后向模式 AD 计算。当前，"forward-mode" 需要 vectorized=True。默认为 "reverse-mode"。如果 func 的输出多于输入，则 "forward-mode" 通常性能更好。否则，倾向于使用 "reverse-mode"。

返回

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

返回类型

雅可比矩阵 (Tensor 或嵌套元组形式的 Tensors)

示例

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

>>> def linear_model(x):
...     W = torch.tensor([[2.0, -1.0], [0.0, 1.0]])
...     b = torch.tensor([1.0, 0.5])
...     return x @ W.T + b

>>> x = torch.randn(4, 2, requires_grad=True)
>>> jac = jacobian(linear_model, x, vectorize=True)
>>> jac.shape
torch.Size([4, 2, 4, 2])