torch.Tensor.register_hook¶
- Tensor.register_hook(hook)[source]¶
注册反向钩子。
每次计算关于张量的梯度时,都会调用该钩子。钩子应具有以下签名
hook(grad) -> Tensor or None
钩子不应修改其参数,但可以选择返回一个新梯度,该梯度将用于代替
grad
。此函数返回一个句柄,该句柄具有一个方法
handle.remove()
,用于从模块中删除钩子。注意
有关此钩子何时执行以及其执行相对于其他钩子的顺序,请参阅 反向钩子执行。
示例
>>> v = torch.tensor([0., 0., 0.], requires_grad=True) >>> h = v.register_hook(lambda grad: grad * 2) # double the gradient >>> v.backward(torch.tensor([1., 2., 3.])) >>> v.grad 2 4 6 [torch.FloatTensor of size (3,)] >>> h.remove() # removes the hook