Python PyTorch jacobian用法及代碼示例

本文簡要介紹python語言中 torch.autograd.functional.jacobian 的用法。

用法:

torch.autograd.functional.jacobian(func, inputs, create_graph=False, strict=False, vectorize=False)

參數：

• func(函數) -一個 Python 函數，它接受張量輸入並返回張量元組或張量。

• inputs(張量元組或者Tensor) -函數 func 的輸入。

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

• strict(bool,可選的) -如果 True ，當我們檢測到存在一個輸入使得所有輸出都獨立於它時，將引發錯誤。如果 False ，我們返回一個零張量作為所述輸入的雅可比，這是預期的數學值。默認為 False 。

• vectorize(bool,可選的) -此函數是實驗性的，請自行承擔使用風險。在計算雅可比時，通常我們在雅可比的每一行調用一次autograd.grad。如果此標誌是 True ，我們使用 vmap 原型函數作為後端來矢量化對 autograd.grad 的調用，因此我們隻調用它一次而不是每行一次。這應該會在許多用例中帶來性能改進，但是，由於此函數不完整，可能會出現性能懸崖。請使用torch._C._debug_only_display_vmap_fallback_warnings(True) 顯示任何性能警告，如果您的用例存在警告，請向我們提交問題。默認為 False 。

返回：

如果隻有一個輸入和輸出，這將是一個包含線性化輸入和輸出的雅可比行列式的張量。如果兩者中的一個是元組，那麽雅可比將是張量的元組。如果它們都是元組，那麽雅可比將是張量元組的元組，其中 Jacobian[i][j] 將包含第 i 輸出和 j 輸入的雅可比，並且將具有大小的串聯相應的輸出和相應的輸入將具有與相應輸入相同的 dtype 和設備。

返回類型：

雅可比行列式(Tensor 或張量的嵌套元組)

計算給定函數的雅可比行列式的函數。

示例

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