以pytorch的forward hook为例探究hook机制

在看pytorch的nn.Module部分的源码的时候，看到了一堆"钩子"，也就是hook，然后去研究了一下这是啥玩意。
基本概念

在深度学习中，hook 是一种可以在模型的不同阶段插入自界说代码的机制。通过自界说数据在通过模型的特定层的额外举动，可以用来监控状态，帮忙调试，获得中间结果。
以前向hook为例

前向hook是模型在forward过程中会调用的hook，通过torch.nn.Module的register_forward_hook() 函数，将一个自界说的hook函数注册给模型的一个层
该层在进行前向之后，根据其输入和输出会进行相应的举动。

1. import torch
2. import torch.nn as nn
4. # 定义模型
5. class SimpleModel(nn.Module):
6.     def __init__(self):
7.         super(SimpleModel, self).__init__()
8.         self.fc1 = nn.Linear(10, 5)
9.         self.fc2 = nn.Linear(5, 2)
11.     def forward(self, x):
12.         x = self.fc1(x)
13.         return self.fc2(x)
15. model = SimpleModel()
17. # 自定义的forward hook
18. def my_forward_hook(module, input, output):
19.     print(f"层: {module}")
20.     print(f"输入: {input}")
21.     print(f"输出: {output}")
24. # 为模型的fc1层注册hook
25. hook = model.fc1.register_forward_hook(my_forward_hook)
27. # 移除这个hook
28. hook.remove()

复制代码

接口

hook函数的格式

须要是一个担当三个特定参数，返回None的函数

1. def hook_function(module, input, output):
2.     # 自定义逻辑
3.     return None

复制代码

• module: 触发钩子的模型层，事实上是调用register_forward_hook的nn.Module实例
  
• input: 传递给该层的输入张量（大概是元组），是前向流传时该层接收到的输入。
  
• output: 该层的输出张量，是前向流传时该层天生的输出。
  函数内部可以做自界说举动，可以在函数内部对output进行修改，从而改变模型的输出。
  

注册hook

1. hook = model.fc1.register_forward_hook(my_forward_hook)
3. hook.remove()

复制代码

对于界说好的hook函数，将其作为参数，调用须要注册的模型层的注册函数即可。
如果不再须要这个hook，调用remove函数。
简朴的源码讨论

还是以forward hook为例。一个nn.Module具有成员_forward_hooks，这是一个有序字典，在__init__()函数调用的时候被初始化

1. self._forward_hooks = OrderedDict()

复制代码

注册钩子的register函数。
每个hook对应一个RemovableHandle对象，以其id作为键注册到hook字典中，使用其remove函数实现移除。

1. def register_forward_hook(
2.         self,
3.         hook: Union[
4.             Callable[[T, Tuple[Any, ...], Any], Optional[Any]],
5.             Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
6.         ],
7.         *,
8.         prepend: bool = False,
9.         with_kwargs: bool = False,
10.         always_call: bool = False,
11.     ) -> RemovableHandle:
12.         handle = RemovableHandle(
13.             self._forward_hooks,
14.             extra_dict=[
15.                 self._forward_hooks_with_kwargs,
16.                 self._forward_hooks_always_called,
17.             ],
18.         )
19.         self._forward_hooks[handle.id] = hook
20.         if with_kwargs:
21.             self._forward_hooks_with_kwargs[handle.id] = True
22.         if always_call:
23.             self._forward_hooks_always_called[handle.id] = True
24.         if prepend:
25.             self._forward_hooks.move_to_end(handle.id, last=False)  # type: ignore[attr-defined]
26.         return handle
28. #简化版的handle类
29. class RemovableHandle:
30.     def __init__(self, hooks_dict, handle_id):
31.         self.hooks_dict = hooks_dict
32.         self.id = handle_id
34.     def remove(self):
35.         del self.hooks_dict[self.id]

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。