快捷方式

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 (intTuple[int]) – 可选,整数或整数元组,表示获取关于哪些参数的雅可比矩阵。默认值:0。

  • has_aux (bool) – 标志,指示 func 返回一个 (output, aux) 元组,其中第一个元素是需要微分的函数的输出,第二个元素是不会被微分的辅助对象。默认值:False。

  • chunk_size (Noneint) – 如果为 None(默认值),则使用最大块大小(等效于对 vjp 进行一次 vmap 以计算雅可比矩阵)。如果为 1,则使用 for 循环逐行计算雅可比矩阵。如果非 None,则每次计算 chunk_size 行的雅可比矩阵(等效于对 vjp 进行多次 vmap)。如果您在计算雅可比矩阵时遇到内存问题,请尝试指定一个非 None 的 chunk_size。

返回值

返回一个函数,该函数接受与 func 相同的输入,并返回 func 关于 argnums 处的参数的雅可比矩阵。如果 has_aux is True,则返回的函数改为返回一个 (jacobian, aux) 元组,其中 jacobian 是雅可比矩阵,auxfunc 返回的辅助对象。

使用逐点一元运算的基本用法将得到一个对角阵作为雅可比矩阵。

>>> 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_gradjacrev 一起使用。情况 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 之外的上下文管理器的结果。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取适合初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源