torch.autograd.graph.Node.register_prehook¶
- abstract Node.register_prehook(fn)[原始碼][原始碼]¶
註冊一個 backward pre-hook。
每次計算相對於 Node 的梯度時,都會呼叫這個 hook。這個 hook 應該具有以下簽名:
hook(grad_outputs: Tuple[Tensor]) -> Tuple[Tensor] or None
這個 hook 不應修改其參數,但它可以選擇性地傳回一個新的梯度,該梯度將用於代替
grad_outputs
。此函式傳回一個帶有
handle.remove()
方法的控制代碼,該方法會從模組中移除該 hook。注意
有關何時執行此 hook 以及其執行順序相對於其他 hook 的更多資訊,請參閱Backward Hooks execution。
範例
>>> 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