torch.Tensor.register_hook

Tensor.register_hook(hook)

注册一个反向钩子。

每次计算关于 Tensor 的梯度时都会调用该 hook。该 hook 应具有以下签名：

hook(grad) -> Tensor or None

该 hook 不应修改其参数，但可以选择性地返回一个新的梯度，该梯度将用于替换 grad。

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

注意

有关此 hook 何时执行以及如何与其他 hook 排序执行的信息，请参阅 反向传播 hook 的执行。

示例

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