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