torch.autograd.graph.Node.register_prehook

abstract Node.register_prehook(fn)

注册一个反向传播前置钩子 (pre-hook)。

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

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

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

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

注意

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

示例

>>> 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_prehook(lambda gI: (gI[0] * 2,))
>>> b.sum().backward(retain_graph=True)
>>> print(a.grad)
tensor([2., 2., 2.])
>>> handle.remove()
>>> a.grad = None
>>> b.sum().backward(retain_graph=True)
>>> print(a.grad)
tensor([1., 1., 1.])

返回类型

RemovableHandle (可移除句柄)