以pytorch的forward hook为例探究hook机制
在看pytorch的nn.Module部分的源码的时候,看到了一堆"钩子",也就是hook,然后去研究了一下这是啥玩意。
基本概念
在深度学习中,hook 是一种可以在模型的不同阶段插入自定义代码的机制。通过自定义数据在通过模型的特定层的额外行为,可以用来监控状态,协助调试,获得中间结果。
以前向hook为例
前向hook是模型在forward过程中会调用的hook,通过torch.nn.Module的register_forward_hook() 函数,将一个自定义的hook函数注册给模型的一个层
该层在进行前向之后,根据其输入和输出会进行相应的行为。
import torch
import torch.nn as nn
# 定义模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)
def forward(self, x):
x = self.fc1(x)
return self.fc2(x)
model = SimpleModel()
# 自定义的forward hook
def my_forward_hook(module, input, output):
print(f"层: {module}")
print(f"输入: {input}")
print(f"输出: {output}")
# 为模型的fc1层注册hook
hook = model.fc1.register_forward_hook(my_forward_hook)
# 移除这个hook
hook.remove()
接口
hook函数的格式
需要是一个接受三个特定参数,返回None的函数
def hook_function(module, input, output):
# 自定义逻辑
return None
- module: 触发钩子的模型层,事实上是调用register_forward_hook的nn.Module实例
- input: 传递给该层的输入张量(可能是元组),是前向传播时该层接收到的输入。
- output: 该层的输出张量,是前向传播时该层生成的输出。
函数内部可以做自定义行为,可以在函数内部对output进行修改,从而改变模型的输出。
注册hook
hook = model.fc1.register_forward_hook(my_forward_hook)
hook.remove()
对于定义好的hook函数,将其作为参数,调用需要注册的模型层的注册函数即可。
如果不再需要这个hook,调用remove函数。
简单的源码讨论
还是以forward hook为例。一个nn.Module具有成员_forward_hooks,这是一个有序字典,在__init__()函数调用的时候被初始化
self._forward_hooks = OrderedDict()
注册钩子的register函数。
每个hook对应一个RemovableHandle对象,以其id作为键注册到hook字典中,利用其remove函数实现移除。
def register_forward_hook(
self,
hook: Union[
Callable[[T, Tuple[Any, ...], Any], Optional[Any]],
Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
],
*,
prepend: bool = False,
with_kwargs: bool = False,
always_call: bool = False,
) -> RemovableHandle:
handle = RemovableHandle(
self._forward_hooks,
extra_dict=[
self._forward_hooks_with_kwargs,
self._forward_hooks_always_called,
],
)
self._forward_hooks[handle.id] = hook
if with_kwargs:
self._forward_hooks_with_kwargs[handle.id] = True
if always_call:
self._forward_hooks_always_called[handle.id] = True
if prepend:
self._forward_hooks.move_to_end(handle.id, last=False) # type: ignore[attr-defined]
return handle
#简化版的handle类
class RemovableHandle:
def __init__(self, hooks_dict, handle_id):
self.hooks_dict = hooks_dict
self.id = handle_id
def remove(self):
del self.hooks_dict[self.id]