torch.Tensor.register_hook

Tensor.register_hook(hook)

注册一个反向钩子。

每次计算关于该张量的梯度时，都会调用该钩子。钩子应具有以下签名：

hook(grad) -> Tensor 或 无

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

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

注意

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

示例：

>>> v = torch.tensor([0., 0., 0.], requires_grad=True)
>>> h = v.register_hook(lambda grad: grad * 2)  # 将梯度加倍
>>> v.backward(torch.tensor([1., 2., 3.]))
>>> v.grad

 2
 4
 6
[torch.FloatTensor of size (3,)]

>>> h.remove()  # 移除钩子