torch.autograd.graph.Node.register_hook

抽象 Node.register_hook(fn)

注册一个反向钩子。

每当计算相对于该 Node 的梯度时，就会调用此钩子。该钩子应具有以下签名

hook(grad_inputs: Tuple[Tensor], grad_outputs: Tuple[Tensor]) -> Tuple[Tensor] or None

钩子不应修改其参数，但可以选择返回一个新的梯度，该梯度将取代 grad_inputs 使用。

此函数返回一个句柄，该句柄有一个 handle.remove() 方法，用于从模块中移除钩子。

注意

有关此钩子何时执行以及其执行顺序相对于其他钩子的更多信息，请参阅反向钩子执行。

注意

在极少数情况下，如果在 Node 已开始执行时注册钩子，则无法保证 grad_outputs 的内容（它可能像往常一样或为空，取决于其他因素）。钩子仍然可以选择返回一个新的梯度来代替 grad_inputs 使用，而与 grad_outputs 无关。

示例

>>> import torch
>>> a = torch.tensor([0., 0., 0.], requires_grad=True)
>>> b = a.clone()
>>> assert isinstance(b.grad_fn, torch.autograd.graph.Node)
>>> handle = b.grad_fn.register_hook(lambda gI, gO: (gO[0] * 2,))
>>> b.sum().backward(retain_graph=True)
>>> print(a.grad)
tensor([2., 2., 2.])
>>> handle.remove() # Removes the hook
>>> a.grad = None
>>> b.sum().backward(retain_graph=True)
>>> print(a.grad)
tensor([1., 1., 1.])

返回类型

RemovableHandle