Python PyTorch Module.register_forward_hook用法及代碼示例

本文簡要介紹python語言中 torch.nn.Module.register_forward_hook 的用法。

用法:

register_forward_hook(hook)

返回：

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

返回類型：

torch.utils.hooks.RemovableHandle

在模塊上注冊一個前向掛鉤。

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

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

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