torch.func.jacrev

torch.func.jacrev(func, argnums=0, *, has_aux=False, chunk_size=None, _preallocate_and_copy=False)

使用反向模式自動微分，計算 func 相對於索引為 argnum 的參數的 Jacobian 矩陣

注意

使用 chunk_size=1 相當於使用 for 迴圈逐行計算 Jacobian 矩陣，也就是說，vmap() 的限制不適用。

參數

• func (function) – 一個 Python 函式，接受一個或多個參數，其中一個必須是 Tensor，並傳回一個或多個 Tensors

• argnums (int 或 Tuple[int]) – 可選，整數或整數元組，指定要計算 Jacobian 矩陣的參數。預設值：0。

• has_aux (bool) – 標記，表示 func 傳回一個 (output, aux) 元組，其中第一個元素是要微分的函式的輸出，第二個元素是不會被微分的輔助物件。預設值：False。

• chunk_size (None 或 int) – 如果為 None (預設值)，則使用最大區塊大小 (相當於對 vjp 進行單一 vmap 以計算 Jacobian 矩陣)。 如果為 1，則使用 for 迴圈逐行計算 Jacobian 矩陣。 如果不為 None，則一次計算 chunk_size 列 Jacobian 矩陣（相當於對 vjp 進行多次 vmap）。 如果在計算 Jacobian 矩陣時遇到記憶體問題，請嘗試指定非 None 的 chunk_size。

傳回值

傳回一個函式，該函式接受與 func 相同的輸入，並傳回 func 相對於 argnums 的參數的 Jacobian 矩陣。 如果 has_aux 為 True，則傳回的函式會傳回一個 (jacobian, aux) 元組，其中 jacobian 是 Jacobian 矩陣，而 aux 是 func 傳回的輔助物件。

使用點式一元運算的基本用法會產生對角陣列作為 Jacobian 矩陣

>>> from torch.func import jacrev
>>> x = torch.randn(5)
>>> jacobian = jacrev(torch.sin)(x)
>>> expected = torch.diag(torch.cos(x))
>>> assert torch.allclose(jacobian, expected)

如果您想要同時計算函式的輸出和函式的 Jacobian 矩陣，請使用 has_aux 標記將輸出作為輔助物件傳回

>>> from torch.func import jacrev
>>> x = torch.randn(5)
>>>
>>> def f(x):
>>>   return x.sin()
>>>
>>> def g(x):
>>>   result = f(x)
>>>   return result, result
>>>
>>> jacobian_f, f_x = jacrev(g, has_aux=True)(x)
>>> assert torch.allclose(f_x, f(x))

jacrev() 可以與 vmap 組合以產生批次 Jacobian 矩陣

>>> from torch.func import jacrev, vmap
>>> x = torch.randn(64, 5)
>>> jacobian = vmap(jacrev(torch.sin))(x)
>>> assert jacobian.shape == (64, 5, 5)

此外，jacrev() 可以與自身組合以產生 Hessian 矩陣

>>> from torch.func import jacrev
>>> def f(x):
>>>   return x.sin().sum()
>>>
>>> x = torch.randn(5)
>>> hessian = jacrev(jacrev(f))(x)
>>> assert torch.allclose(hessian, torch.diag(-x.sin()))

預設情況下，jacrev() 計算相對於第一個輸入的 Jacobian 矩陣。 但是，它可以通過使用 argnums 來計算相對於不同參數的 Jacboian 矩陣

>>> from torch.func import jacrev
>>> def f(x, y):
>>>   return x + y ** 2
>>>
>>> x, y = torch.randn(5), torch.randn(5)
>>> jacobian = jacrev(f, argnums=1)(x, y)
>>> expected = torch.diag(2 * y)
>>> assert torch.allclose(jacobian, expected)

此外，將元組傳遞給 argnums 將會計算相對於多個參數的 Jacobian 矩陣

>>> from torch.func import jacrev
>>> def f(x, y):
>>>   return x + y ** 2
>>>
>>> x, y = torch.randn(5), torch.randn(5)
>>> jacobian = jacrev(f, argnums=(0, 1))(x, y)
>>> expectedX = torch.diag(torch.ones_like(x))
>>> expectedY = torch.diag(2 * y)
>>> assert torch.allclose(jacobian[0], expectedX)
>>> assert torch.allclose(jacobian[1], expectedY)

注意

將 PyTorch torch.no_grad 與 jacrev 一起使用。 情況 1：在函式內部使用 torch.no_grad

>>> def f(x):
>>>     with torch.no_grad():
>>>         c = x ** 2
>>>     return x - c

在這種情況下，jacrev(f)(x) 將遵守內部的 torch.no_grad。

情況 2：在 torch.no_grad 上下文管理器內部使用 jacrev

>>> with torch.no_grad():
>>>     jacrev(f)(x)

在這種情況下，jacrev 將遵守內部的 torch.no_grad，但不遵守外部的 torch.no_grad。 這是因為 jacrev 是一種「函式轉換」：其結果不應取決於 f 外部的上下文管理器的結果。