快捷方式

torch.func.vjp

torch.func.vjp(func, *primals, has_aux=False)

代表向量-雅可比积,返回一个元组,包含应用于 primalsfunc 的结果和一个函数,该函数在给定 cotangents 时,计算 func 相对于 primals 的反向模式雅可比乘以 cotangents

参数
  • func (Callable) – 接受一个或多个参数的 Python 函数。必须返回一个或多个张量。

  • primals (Tensors) – func 的位置参数,必须全部为张量。返回的函数还将计算相对于这些参数的导数

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

返回值

返回一个 (output, vjp_fn) 元组,包含应用于 primalsfunc 的输出和一个函数,该函数使用传递给返回函数的切线计算 func 相对于所有 primals 的 vjp。如果 has_aux is True,则返回一个 (output, vjp_fn, aux) 元组。返回的 vjp_fn 函数将返回每个 VJP 的元组。

在简单情况下使用时,vjp() 的行为与 grad() 相同

>>> x = torch.randn([5])
>>> f = lambda x: x.sin().sum()
>>> (_, vjpfunc) = torch.func.vjp(f, x)
>>> grad = vjpfunc(torch.tensor(1.))[0]
>>> assert torch.allclose(grad, torch.func.grad(f)(x))

但是,vjp() 可以通过传递每个输出的切线来支持具有多个输出的函数

>>> x = torch.randn([5])
>>> f = lambda x: (x.sin(), x.cos())
>>> (_, vjpfunc) = torch.func.vjp(f, x)
>>> vjps = vjpfunc((torch.ones([5]), torch.ones([5])))
>>> assert torch.allclose(vjps[0], x.cos() + -x.sin())

vjp() 甚至可以支持输出为 Python 结构体

>>> x = torch.randn([5])
>>> f = lambda x: {'first': x.sin(), 'second': x.cos()}
>>> (_, vjpfunc) = torch.func.vjp(f, x)
>>> cotangents = {'first': torch.ones([5]), 'second': torch.ones([5])}
>>> vjps = vjpfunc(cotangents)
>>> assert torch.allclose(vjps[0], x.cos() + -x.sin())

vjp() 返回的函数将计算相对于每个 primals 的偏导数

>>> x, y = torch.randn([5, 4]), torch.randn([4, 5])
>>> (_, vjpfunc) = torch.func.vjp(torch.matmul, x, y)
>>> cotangents = torch.randn([5, 5])
>>> vjps = vjpfunc(cotangents)
>>> assert len(vjps) == 2
>>> assert torch.allclose(vjps[0], torch.matmul(cotangents, y.transpose(0, 1)))
>>> assert torch.allclose(vjps[1], torch.matmul(x.transpose(0, 1), cotangents))

primalsf 的位置参数。所有关键字参数都使用其默认值

>>> x = torch.randn([5])
>>> def f(x, scale=4.):
>>>   return x * scale
>>>
>>> (_, vjpfunc) = torch.func.vjp(f, x)
>>> vjps = vjpfunc(torch.ones_like(x))
>>> assert torch.allclose(vjps[0], torch.full(x.shape, 4.))

注意

将 PyTorch torch.no_gradvjp 结合使用。情况 1:在函数内部使用 torch.no_grad

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

在这种情况下,vjp(f)(x) 将尊重内部的 torch.no_grad

情况 2:在 torch.no_grad 上下文管理器内部使用 vjp

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

在这种情况下,vjp 将尊重内部的 torch.no_grad,但不会尊重外部的 torch.no_grad。这是因为 vjp 是一种“函数转换”:其结果不应取决于 f 外部的上下文管理器的结果。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

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

查看资源