torch.autograd.graph.Node.register_prehook
- 抽象节点.注册前置钩子(fn)[源代码]
-
注册一个逆向前置钩子。
每次计算与节点相关的梯度时,都会触发此钩子函数。该钩子函数应具有以下签名:
hook(grad_outputs: Tuple[Tensor]) -> Tuple[Tensor] or None
钩子函数不应修改其输入参数,但可以可选地返回一个新的梯度来替换
grad_outputs
。此函数返回一个包含
handle.remove()
方法的句柄,该方法可以用于从模块中移除钩子。注意
关于此钩子的执行时机及其与其它钩子执行顺序的关系,请参见Backward Hooks 执行。
示例:
>>> 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.])
- 返回类型
-
可移除句柄