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

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

用法:

register_forward_pre_hook(hook)

返回：

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

返回類型：

torch.utils.hooks.RemovableHandle

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

每次調用 forward() 之前都會調用該鉤子。它應該具有以下簽名：

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

輸入僅包含給模塊的位置參數。關鍵字參數不會傳遞給鉤子，隻會傳遞給 forward 。鉤子可以修改輸入。用戶可以在鉤子中返回一個元組或單個修改值。如果返回單個值，我們會將值包裝到一個元組中(除非該值已經是一個元組)。