torch.func.grad¶
- torch.func.grad(func, argnums=0, has_aux=False)¶
grad
运算符有助于计算func
相对于由argnums
指定的输入的梯度。此运算符可以嵌套以计算高阶梯度。- 参数
- 返回
用于计算其输入梯度的函数。默认情况下,函数的输出是相对于第一个参数的梯度张量。如果指定的
has_aux
等于True
,则返回梯度和输出辅助对象的元组。如果argnums
是一个整数元组,则返回一个相对于每个argnums
值的输出梯度元组。- 返回类型
使用
grad
的示例>>> from torch.func import grad >>> x = torch.randn([]) >>> cos_x = grad(lambda x: torch.sin(x))(x) >>> assert torch.allclose(cos_x, x.cos()) >>> >>> # Second-order gradients >>> neg_sin_x = grad(grad(lambda x: torch.sin(x)))(x) >>> assert torch.allclose(neg_sin_x, -x.sin())
当与
vmap
组合使用时,grad
可用于计算每个样本的梯度>>> from torch.func import grad, vmap >>> batch_size, feature_size = 3, 5 >>> >>> def model(weights, feature_vec): >>> # Very simple linear model with activation >>> assert feature_vec.dim() == 1 >>> return feature_vec.dot(weights).relu() >>> >>> def compute_loss(weights, example, target): >>> y = model(weights, example) >>> return ((y - target) ** 2).mean() # MSELoss >>> >>> weights = torch.randn(feature_size, requires_grad=True) >>> examples = torch.randn(batch_size, feature_size) >>> targets = torch.randn(batch_size) >>> inputs = (weights, examples, targets) >>> grad_weight_per_example = vmap(grad(compute_loss), in_dims=(None, 0, 0))(*inputs)
使用
grad
以及has_aux
和argnums
的示例>>> from torch.func import grad >>> def my_loss_func(y, y_pred): >>> loss_per_sample = (0.5 * y_pred - y) ** 2 >>> loss = loss_per_sample.mean() >>> return loss, (y_pred, loss_per_sample) >>> >>> fn = grad(my_loss_func, argnums=(0, 1), has_aux=True) >>> y_true = torch.rand(4) >>> y_preds = torch.rand(4, requires_grad=True) >>> out = fn(y_true, y_preds) >>> # > output is ((grads w.r.t y_true, grads w.r.t y_preds), (y_pred, loss_per_sample))
注意
将 PyTorch
torch.no_grad
与grad
结合使用。情况 1:在函数内部使用
torch.no_grad
>>> def f(x): >>> with torch.no_grad(): >>> c = x ** 2 >>> return x - c
在这种情况下,
grad(f)(x)
将遵守内部的torch.no_grad
。情况 2:在
torch.no_grad
上下文管理器内部使用grad
>>> with torch.no_grad(): >>> grad(f)(x)
在这种情况下,
grad
将遵守内部的torch.no_grad
,但不遵守外部的。这是因为grad
是一个“函数转换”:其结果不应依赖于f
外部的上下文管理器的结果。