torch.autograd.graph.Node.register_prehook#
- abstract Node.register_prehook(fn)[source]#
註冊一個後向 pre-hook。
每次計算 Node 的梯度時都會呼叫此 hook。hook 應該具有以下簽名:
hook(grad_outputs: Tuple[Tensor]) -> Tuple[Tensor] or None
hook 不應該修改其引數,但可以選擇返回一個新的梯度,該梯度將用於替換
grad_outputs。此函式返回一個控制代碼,其中包含一個方法
handle.remove(),用於從模組中移除該鉤子。注意
有關此 hook 何時執行以及如何與其他 hook 排序執行的資訊,請參閱 反向傳播 hook 的執行。
示例
>>> 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_prehook(lambda gI: (gI[0] * 2,)) >>> b.sum().backward(retain_graph=True) >>> print(a.grad) tensor([2., 2., 2.]) >>> handle.remove() >>> a.grad = None >>> b.sum().backward(retain_graph=True) >>> print(a.grad) tensor([1., 1., 1.])
- 返回型別
RemovableHandle