torch.nn.modules.module.register_module_forward_hook#
- torch.nn.modules.module.register_module_forward_hook(hook, *, with_kwargs=False, always_call=False)[source]#
为所有模块注册一个全局前向钩子。
警告
此函数会向 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
注册)之前执行。