Python PyTorch ScriptModule.register_full_backward_hook用法及代碼示例

本文簡要介紹python語言中 torch.jit.ScriptModule.register_full_backward_hook 的用法。

用法:

register_full_backward_hook(hook)

返回：

一個句柄，可用於通過調用 handle.remove() 刪除添加的鉤子

返回類型：

torch.utils.hooks.RemovableHandle

在模塊上注冊一個後向鉤子。

每次計算與模塊輸入相關的梯度時，都會調用該鉤子。掛鉤應具有以下簽名：

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

grad_input 和 grad_output 是元組，分別包含關於輸入和輸出的梯度。鉤子不應該修改它的參數，但它可以選擇返回一個關於輸入的新梯度，該梯度將在後續計算中用於代替grad_input。 grad_input 將僅對應於作為位置參數給出的輸入，並且所有 kwarg 參數都被忽略。對於所有非張量參數，grad_input 和 grad_output 中的條目將為 None。

由於技術原因，當這個鉤子被應用到一個 Module 時，它的 forward 函數將接收傳遞給 Module 的每個 Tensor 的視圖。同樣，調用者將收到模塊 forward 函數返回的每個張量的視圖。

警告

使用後向鉤子時不允許就地修改輸入或輸出，這會引發錯誤。