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 (Tensors 元組或Tensor) – 函式 func 的輸入。

• create_graph (bool, 可選) – 如果為 True，雅可比矩陣將以可微分的方式計算。請注意，當 strict 為 False 時，結果可能不需要梯度或與輸入斷開連線。預設為 False。

• strict (bool, 可選) – 如果為 True，當我們檢測到存在某個輸入，所有輸出都與其無關時，將引發錯誤。如果為 False，則對於這些輸入，我們將返回一個全零 Tensor 作為雅可比矩陣，這是預期的數學值。預設為 False。

• vectorize (bool, 可選) – 此特性為實驗性。如果您尋求更穩定且效能更好的功能，請考慮改用torch.func.jacrev() 或 torch.func.jacfwd()。在計算雅可比矩陣時，通常每計算雅可比矩陣的一行都會呼叫一次 autograd.grad。如果此標誌為 True，我們將僅執行一次 autograd.grad 呼叫，並設定 batched_grad=True，它使用了 vmap 原型特性。雖然這在許多情況下會帶來效能提升，但由於此特性仍處於實驗階段，可能存在效能瓶頸。更多資訊請參見 torch.autograd.grad() 的 batched_grad 引數。

• strategy (str, 可選) – 設定為 "forward-mode" 或 "reverse-mode"，以確定使用前向模式 AD 還是反向模式 AD 計算雅可比矩陣。目前，"forward-mode" 需要 vectorized=True。預設為 "reverse-mode"。如果 func 的輸出多於輸入，"forward-mode" 傾向於效能更優。否則，優先使用 "reverse-mode"。

返回值

如果只有一個輸入和一個輸出，這將是一個包含線性化輸入和輸出的雅可比矩陣的單個 Tensor。如果其中一個為元組，則雅可比矩陣將是 Tensors 的元組。如果兩者都是元組，則雅可比矩陣將是 Tensors 元組的元組，其中 Jacobian[i][j] 包含第 i 個輸出和第 j 個輸入的雅可比矩陣，其大小將是對應輸出和對應輸入大小的拼接，並且與對應輸入具有相同的 dtype 和 device。如果 strategy 為 forward-mode，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.]]))