torch.nn.modules.module.register_module_forward_hook

torch.nn.modules.module.register_module_forward_hook(hook, *, with_kwargs=False, always_call=False)

為所有模組註冊一個全域性前向鉤子。

警告

這會向 nn.module 模組新增全域性狀態，僅用於除錯/效能分析目的。

該鉤子會在 forward() 計算出輸出後每次被呼叫。它應該具有以下簽名

hook(module, input, output) -> None or modified output

輸入只包含傳遞給模組的位置引數。關鍵字引數不會傳遞給鉤子，只會傳遞給 forward。您可以透過返回一個新值來選擇性地修改模組的輸出，該新值將替換 forward() 函式的輸出。

引數

• hook (Callable) – 要註冊的使用者定義鉤子。

• always_call (bool) – 如果為 True，則無論呼叫模組時是否丟擲異常，都會執行該 hook。預設值：False

返回值

一個控制代碼，透過呼叫 handle.remove() 可以移除新增的鉤子。

返回型別

torch.utils.hooks.RemovableHandle

該鉤子將在使用 register_forward_hook 註冊的特定模組鉤子之前執行。