Pytorch 容器 - 4. Module中hook注册：register_forward_pre_hook()，register_forward_hook()

目录

1. hook作用

2. register_forward_pre_hook(hook)

3. register_forward_hook()

4. 代码示例

5. 其他示例

---

1. hook作用

hook是一个可调用的对象，它预定义了函数声明（即函数参数，返回值，调用方式等)。当调用forward() / backward()时，module对应的输入输出都会传到hook上，并可以在hook中处理这些输入输出。因此hook可中进行一些如：可视化中间特征、冻结部分层的等操作

2. register_forward_pre_hook(hook)

• 该函数在foward()之前运行
• 该函数能够修改输入并将修改后的新的输入结果返回给 forward()
• 如果想要移除hook函数可以使用 remove()

3. register_forward_hook()

• 该函数在foward()之后运行
• 该函数能够修改输出结果 (inplace)
• 如果想要移除hook函数可以使用 remove()