torch.autograd.graph.Node.register_hook

抽象节点.注册钩子(fn) [源代码]

注册一个后向钩子。

每次计算与节点相关的梯度时，都会触发此钩子函数。该钩子函数应具有以下签名：

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

钩子函数不应修改其参数，但可以可选地返回一个新的梯度来替换grad_inputs。

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

注意

关于此钩子的执行时机及其与其它钩子执行顺序的关系，请参见Backward Hooks 执行。

注意

在罕见的情况下，如果钩子注册时节点已经开始执行，则无法再保证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.])

返回类型

可移除句柄