torch.func.vjp¶
- torch.func.vjp(func, *primals, has_aux=False)¶
代表向量-雅可比积,返回一个元组,包含应用于
primals
的func
的结果和一个函数,该函数在给定cotangents
时,计算func
相对于primals
的反向模式雅可比乘以cotangents
。- 参数
func (Callable) – 接受一个或多个参数的 Python 函数。必须返回一个或多个张量。
primals (Tensors) –
func
的位置参数,必须全部为张量。返回的函数还将计算相对于这些参数的导数has_aux (bool) – 标志,指示
func
返回一个(output, aux)
元组,其中第一个元素是将要微分的函数的输出,第二个元素是不会进行微分的其他辅助对象。默认值:False。
- 返回值
返回一个
(output, vjp_fn)
元组,包含应用于primals
的func
的输出和一个函数,该函数使用传递给返回函数的切线计算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))
primals
是f
的位置参数。所有关键字参数都使用其默认值>>> 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_grad
与vjp
结合使用。情况 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
外部的上下文管理器的结果。