Python PyTorch register_module_forward_hook用法及代碼示例

本文簡要介紹python語言中 torch.nn.modules.module.register_module_forward_hook 的用法。

用法:

torch.nn.modules.module.register_module_forward_hook(hook)

返回：

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

返回類型：

torch.utils.hooks.RemovableHandle

為所有模塊注冊一個全局前向鉤子

警告

這會將全局狀態添加到 nn.module 模塊，並且僅用於調試/分析目的。

每次 forward() 計算輸出後都會調用該鉤子。它應該具有以下簽名：

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

輸入僅包含給模塊的位置參數。關鍵字參數不會傳遞給鉤子，隻會傳遞給 forward 。鉤子可以修改輸出。它可以就地修改輸入，但不會對轉發產生影響，因為這是在調用 forward() 之後調用的。

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