torch.autograd.graph.Node.register_hook

抽象 Node.register_hook(fn)

註冊一個反向鉤子。

每當計算相對於該 Node 的梯度時，就會呼叫此鉤子。該鉤子應具有以下簽名

hook(grad_inputs: Tuple[Tensor], grad_outputs: Tuple[Tensor]) -> Tuple[Tensor] or None

鉤子不應修改其引數，但可以選擇返回一個新的梯度，該梯度將取代 grad_inputs 使用。

此函式返回一個控制代碼，該控制代碼有一個 handle.remove() 方法，用於從模組中移除鉤子。

注意

有關此鉤子何時執行以及其執行順序相對於其他鉤子的更多資訊，請參閱反向鉤子執行。

注意

在極少數情況下，如果在 Node 已開始執行時註冊鉤子，則無法保證 grad_outputs 的內容（它可能像往常一樣或為空，取決於其他因素）。鉤子仍然可以選擇返回一個新的梯度來代替 grad_inputs 使用，而與 grad_outputs 無關。

示例

>>> import torch
>>> 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_hook(lambda gI, gO: (gO[0] * 2,))
>>> b.sum().backward(retain_graph=True)
>>> print(a.grad)
tensor([2., 2., 2.])
>>> handle.remove() # Removes the hook
>>> a.grad = None
>>> b.sum().backward(retain_graph=True)
>>> print(a.grad)
tensor([1., 1., 1.])

返回型別

RemovableHandle