torch.autograd.graph.Node.register_hook

abstract Node.register_hook(fn)

註冊一個反向傳播 hook。

每次計算相對於 Node 的梯度時，都會調用該 hook。這個 hook 應該具有以下簽章：

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

該 hook 不應修改其參數，但它可以選擇性地返回一個新的梯度，該梯度將用於代替 grad_inputs。

此函數返回一個具有 handle.remove() 方法的 handle，該方法從模組中移除該 hook。

注意

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

注意

在極少數情況下，如果在 Node 已經開始執行時註冊了該 hook，則不再保證 grad_outputs 的內容（它可能像往常一樣，也可能為空，具體取決於其他因素）。 該 hook 仍然可以選擇性地返回一個新的梯度，以代替 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