torch.func.jacrev¶
- torch.func.jacrev(func, argnums=0, *, has_aux=False, chunk_size=None, _preallocate_and_copy=False)¶
使用反向模式自动微分计算
func
关于索引argnum
处的参数的雅可比矩阵。注意
使用
chunk_size=1
等效于使用 for 循环逐行计算雅可比矩阵,即vmap()
的约束不适用。- 参数
func (function) – 一个 Python 函数,它接受一个或多个参数,其中一个必须是张量,并返回一个或多个张量。
argnums (int 或 Tuple[int]) – 可选,整数或整数元组,表示获取关于哪些参数的雅可比矩阵。默认值:0。
has_aux (bool) – 标志,指示
func
返回一个(output, aux)
元组,其中第一个元素是需要微分的函数的输出,第二个元素是不会被微分的辅助对象。默认值:False。chunk_size (None 或 int) – 如果为 None(默认值),则使用最大块大小(等效于对 vjp 进行一次 vmap 以计算雅可比矩阵)。如果为 1,则使用 for 循环逐行计算雅可比矩阵。如果非 None,则每次计算
chunk_size
行的雅可比矩阵(等效于对 vjp 进行多次 vmap)。如果您在计算雅可比矩阵时遇到内存问题,请尝试指定一个非 None 的 chunk_size。
- 返回值
返回一个函数,该函数接受与
func
相同的输入,并返回func
关于argnums
处的参数的雅可比矩阵。如果has_aux is True
,则返回的函数改为返回一个(jacobian, aux)
元组,其中jacobian
是雅可比矩阵,aux
是func
返回的辅助对象。
使用逐点一元运算的基本用法将得到一个对角阵作为雅可比矩阵。
>>> 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)
如果您希望计算函数的输出以及函数的雅可比矩阵,请使用
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 组合以生成批处理雅可比矩阵。>>> 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()
计算关于第一个输入的雅可比矩阵。但是,它可以通过使用argnums
来计算关于不同参数的雅可比矩阵。>>> 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
将计算关于多个参数的雅可比矩阵。>>> 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
,但不尊重外部的。这是因为jacrev
是一个“函数转换”:其结果不应依赖于f
之外的上下文管理器的结果。