torch.Tensor.register_hook

Tensor.register_hook(hook)[源代码]

注册一个后向钩子。

每次计算出相对于张量的梯度时，都会触发此钩子函数。该钩子函数应具有以下签名：

hook(grad) -> Tensor or None

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

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

注意

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

示例：

>>> 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