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，则无论调用 Module 时是否发生异常，都将运行 hook。默认为 False。

返回

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型

torch.utils.hooks.RemovableHandle

此钩子将在已使用 register_forward_hook 注册的特定模块钩子之前执行。