5 分钟掌握 Python 中的 Hook 钩子函数

1. 什么是Hook

经常会听到钩子函数(hook function)这个概念，最近在看目标检测开源框架mmdetection，里面也出现大量Hook的编程方式，那到底什么是hook？hook的作用是什么？

what is hook ？钩子hook，顾名思义，可以理解是一个挂钩，作用是有需要的时候挂一个东西上去。具体的解释是：钩子函数是把我们自己实现的hook函数在某一时刻挂接到目标挂载点上。
hook函数的作用举个例子，hook的概念在windows桌面软件开发很常见，特别是各种事件触发的机制; 比如C++的MFC程序中，要监听鼠标左键按下的时间，MFC提供了一个onLeftKeyDown的钩子函数。很显然，MFC框架并没有为我们实现onLeftKeyDown具体的操作，只是为我们提供一个钩子，当我们需要处理的时候，只要去重写这个函数，把我们需要操作挂载在这个钩子里，如果我们不挂载，MFC事件触发机制中执行的就是空的操作。

从上面可知

hook函数是程序中预定义好的函数，这个函数处于原有程序流程当中（暴露一个钩子出来）
我们需要再在有流程中钩子定义的函数块中实现某个具体的细节，需要把我们的实现，挂接或者注册（register）到钩子里，使得hook函数对目标可用
hook 是一种编程机制，和具体的语言没有直接的关系
如果从设计模式上看，hook模式是模板方法的扩展
钩子只有注册的时候，才会使用，所以原有程序的流程中，没有注册或挂载时，执行的是空（即没有执行任何操作）

本文用python来解释hook的实现方式，并展示在开源项目中hook的应用案例。hook函数和我们常听到另外一个名称：回调函数（callback function）功能是类似的，可以按照同种模式来理解。

2. hook实现例子

据我所知，hook函数最常使用在某种流程处理当中。这个流程往往有很多步骤。hook函数常常挂载在这些步骤中，为增加额外的一些操作，提供灵活性。

下面举一个简单的例子，这个例子的目的是实现一个通用往队列中插入内容的功能。流程步骤有2个

需要再插入队列前，对数据进行筛选 input_filter_fn
插入队列 insert_queue

1class ContentStash(object):  
2    """  
3    content stash for online operation  
4    pipeline is  
5    1. input_filter: filter some contents, no use to user  
6    2. insert_queue(redis or other broker): insert useful content to queue  
7    """  
8  
9    def __init__(self):  
10        self.input_filter_fn = None  
11        self.broker = []  
12  
13    def register_input_filter_hook(self, input_filter_fn):  
14        """  
15        register input filter function, parameter is content dict  
16        Args:  
17            input_filter_fn: input filter function  
18  
19        Returns:  
20  
21        """  
22        self.input_filter_fn = input_filter_fn  
23  
24    def insert_queue(self, content):  
25        """  
26        insert content to queue  
27        Args:  
28            content: dict  
29  
30        Returns:  
31  
32        """  
33        self.broker.append(content)  
34  
35    def input_pipeline(self, content, use=False):  
36        """  
37        pipeline of input for content stash  
38        Args:  
39            use: is use, defaul False  
40            content: dict  
41  
42        Returns:  
43  
44        """  
45        if not use:  
46            return  
47  
48        # input filter  
49        if self.input_filter_fn:  
50            _filter = self.input_filter_fn(content)  
51              
52        # insert to queue  
53        if not _filter:  
54            self.insert_queue(content)  
55  
56  
57  
58# test  
59## 实现一个你所需要的钩子实现：比如如果content 包含time就过滤掉，否则插入队列  
60def input_filter_hook(content):  
61    """  
62    test input filter hook  
63    Args:  
64        content: dict  
65  
66    Returns: None or content  
67  
68    """  
69    if content.get('time') is None:  
70        return  
71    else:  
72        return content  
73  
74  
75# 原有程序  
76content = {'filename': 'test.jpg', 'b64_file': "#test", 'data': {"result": "cat", "probility": 0.9}}  
77content_stash = ContentStash('audit', work_dir='')  
78  
79# 挂上钩子函数， 可以有各种不同钩子函数的实现，但是要主要函数输入输出必须保持原有程序中一致，比如这里是content  
80content_stash.register_input_filter_hook(input_filter_hook)  
81  
82# 执行流程  
83content_stash.input_pipeline(content)  
84  
85

3. hook在开源框架中的应用

3.1 keras

在深度学习训练流程中，hook函数体现的淋漓尽致。

一个训练过程（不包括数据准备），会轮询多次训练集，每次称为一个epoch，每个epoch又分为多个batch来训练。流程先后拆解成：

开始训练
训练一个epoch前
训练一个batch前
训练一个batch后
训练一个epoch后
评估验证集
结束训练

这些步骤是穿插在训练一个batch数据的过程中，这些可以理解成是钩子函数，我们可能需要在这些钩子函数中实现一些定制化的东西，比如在训练一个epoch后我们要保存下训练的模型，在结束训练时用最好的模型执行下测试集的效果等等。

keras中是通过各种回调函数来实现钩子hook功能的。这里放一个callback的父类，定制时只要继承这个父类，实现你过关注的钩子就可以了。

1@keras_export('keras.callbacks.Callback')  
2class Callback(object):  
3  """Abstract base class used to build new callbacks.  
4  
5  Attributes:  
6      params: Dict. Training parameters  
7          (eg. verbosity, batch size, number of epochs...).  
8      model: Instance of `keras.models.Model`.  
9          Reference of the model being trained.  
10  
11  The `logs` dictionary that callback methods  
12  take as argument will contain keys for quantities relevant to  
13  the current batch or epoch (see method-specific docstrings).  
14  """  
15  
16  def __init__(self):  
17    self.validation_data = None  # pylint: disable=g-missing-from-attributes  
18    self.model = None  
19    # Whether this Callback should only run on the chief worker in a  
20    # Multi-Worker setting.  
21    # TODO(omalleyt): Make this attr public once solution is stable.  
22    self._chief_worker_only = None  
23    self._supports_tf_logs = False  
24  
25  def set_params(self, params):  
26    self.params = params  
27  
28  def set_model(self, model):  
29    self.model = model  
30  
31 @doc_controls.for_subclass_implementers  
32 @generic_utils.default  
33  def on_batch_begin(self, batch, logs=None):  
34    """A backwards compatibility alias for `on_train_batch_begin`."""  
35  
36 @doc_controls.for_subclass_implementers  
37 @generic_utils.default  
38  def on_batch_end(self, batch, logs=None):  
39    """A backwards compatibility alias for `on_train_batch_end`."""  
40  
41 @doc_controls.for_subclass_implementers  
42  def on_epoch_begin(self, epoch, logs=None):  
43    """Called at the start of an epoch.  
44  
45    Subclasses should override for any actions to run. This function should only  
46    be called during TRAIN mode.  
47  
48    Arguments:  
49        epoch: Integer, index of epoch.  
50        logs: Dict. Currently no data is passed to this argument for this method  
51          but that may change in the future.  
52    """  
53  
54 @doc_controls.for_subclass_implementers  
55  def on_epoch_end(self, epoch, logs=None):  
56    """Called at the end of an epoch.  
57  
58    Subclasses should override for any actions to run. This function should only  
59    be called during TRAIN mode.  
60  
61    Arguments:  
62        epoch: Integer, index of epoch.  
63        logs: Dict, metric results for this training epoch, and for the  
64          validation epoch if validation is performed. Validation result keys  
65          are prefixed with `val_`.  
66    """  
67  
68 @doc_controls.for_subclass_implementers  
69 @generic_utils.default  
70  def on_train_batch_begin(self, batch, logs=None):  
71    """Called at the beginning of a training batch in `fit` methods.  
72  
73    Subclasses should override for any actions to run.  
74  
75    Arguments:  
76        batch: Integer, index of batch within the current epoch.  
77        logs: Dict, contains the return value of `model.train_step`. Typically,  
78          the values of the `Model`'s metrics are returned.  Example:  
79          `{'loss': 0.2, 'accuracy': 0.7}`.  
80    """  
81    # For backwards compatibility.  
82    self.on_batch_begin(batch, logs=logs)  
83  
84 @doc_controls.for_subclass_implementers  
85 @generic_utils.default  
86  def on_train_batch_end(self, batch, logs=None):  
87    """Called at the end of a training batch in `fit` methods.  
88  
89    Subclasses should override for any actions to run.  
90  
91    Arguments:  
92        batch: Integer, index of batch within the current epoch.  
93        logs: Dict. Aggregated metric results up until this batch.  
94    """  
95    # For backwards compatibility.  
96    self.on_batch_end(batch, logs=logs)  
97  
98 @doc_controls.for_subclass_implementers  
99 @generic_utils.default  
100  def on_test_batch_begin(self, batch, logs=None):  
101    """Called at the beginning of a batch in `evaluate` methods.  
102  
103    Also called at the beginning of a validation batch in the `fit`  
104    methods, if validation data is provided.  
105  
106    Subclasses should override for any actions to run.  
107  
108    Arguments:  
109        batch: Integer, index of batch within the current epoch.  
110        logs: Dict, contains the return value of `model.test_step`. Typically,  
111          the values of the `Model`'s metrics are returned.  Example:  
112          `{'loss': 0.2, 'accuracy': 0.7}`.  
113    """  
114  
115 @doc_controls.for_subclass_implementers  
116 @generic_utils.default  
117  def on_test_batch_end(self, batch, logs=None):  
118    """Called at the end of a batch in `evaluate` methods.  
119  
120    Also called at the end of a validation batch in the `fit`  
121    methods, if validation data is provided.  
122  
123    Subclasses should override for any actions to run.  
124  
125    Arguments:  
126        batch: Integer, index of batch within the current epoch.  
127        logs: Dict. Aggregated metric results up until this batch.  
128    """  
129  
130 @doc_controls.for_subclass_implementers  
131 @generic_utils.default  
132  def on_predict_batch_begin(self, batch, logs=None):  
133    """Called at the beginning of a batch in `predict` methods.  
134  
135    Subclasses should override for any actions to run.  
136  
137    Arguments:  
138        batch: Integer, index of batch within the current epoch.  
139        logs: Dict, contains the return value of `model.predict_step`,  
140          it typically returns a dict with a key 'outputs' containing  
141          the model's outputs.  
142    """  
143  
144 @doc_controls.for_subclass_implementers  
145 @generic_utils.default  
146  def on_predict_batch_end(self, batch, logs=None):  
147    """Called at the end of a batch in `predict` methods.  
148  
149    Subclasses should override for any actions to run.  
150  
151    Arguments:  
152        batch: Integer, index of batch within the current epoch.  
153        logs: Dict. Aggregated metric results up until this batch.  
154    """  
155  
156 @doc_controls.for_subclass_implementers  
157  def on_train_begin(self, logs=None):  
158    """Called at the beginning of training.  
159  
160    Subclasses should override for any actions to run.  
161  
162    Arguments:  
163        logs: Dict. Currently no data is passed to this argument for this method  
164          but that may change in the future.  
165    """  
166  
167 @doc_controls.for_subclass_implementers  
168  def on_train_end(self, logs=None):  
169    """Called at the end of training.  
170  
171    Subclasses should override for any actions to run.  
172  
173    Arguments:  
174        logs: Dict. Currently the output of the last call to `on_epoch_end()`  
175          is passed to this argument for this method but that may change in  
176          the future.  
177    """  
178  
179 @doc_controls.for_subclass_implementers  
180  def on_test_begin(self, logs=None):  
181    """Called at the beginning of evaluation or validation.  
182  
183    Subclasses should override for any actions to run.  
184  
185    Arguments:  
186        logs: Dict. Currently no data is passed to this argument for this method  
187          but that may change in the future.  
188    """  
189  
190 @doc_controls.for_subclass_implementers  
191  def on_test_end(self, logs=None):  
192    """Called at the end of evaluation or validation.  
193  
194    Subclasses should override for any actions to run.  
195  
196    Arguments:  
197        logs: Dict. Currently the output of the last call to  
198          `on_test_batch_end()` is passed to this argument for this method  
199          but that may change in the future.  
200    """  
201  
202 @doc_controls.for_subclass_implementers  
203  def on_predict_begin(self, logs=None):  
204    """Called at the beginning of prediction.  
205  
206    Subclasses should override for any actions to run.  
207  
208    Arguments:  
209        logs: Dict. Currently no data is passed to this argument for this method  
210          but that may change in the future.  
211    """  
212  
213 @doc_controls.for_subclass_implementers  
214  def on_predict_end(self, logs=None):  
215    """Called at the end of prediction.  
216  
217    Subclasses should override for any actions to run.  
218  
219    Arguments:  
220        logs: Dict. Currently no data is passed to this argument for this method  
221          but that may change in the future.  
222    """  
223  
224  def _implements_train_batch_hooks(self):  
225    """Determines if this Callback should be called for each train batch."""  
226    return (not generic_utils.is_default(self.on_batch_begin) or  
227            not generic_utils.is_default(self.on_batch_end) or  
228            not generic_utils.is_default(self.on_train_batch_begin) or  
229            not generic_utils.is_default(self.on_train_batch_end))  
230

这些钩子的原始程序是在模型训练流程中的

keras源码位置： tensorflow\python\keras\engine\training.py

部分摘录如下(## I am hook)：

1# Container that configures and calls `tf.keras.Callback`s.  
2      if not isinstance(callbacks, callbacks_module.CallbackList):  
3        callbacks = callbacks_module.CallbackList(  
4            callbacks,  
5            add_history=True,  
6            add_progbar=verbose != 0,  
7            model=self,  
8            verbose=verbose,  
9            epochs=epochs,  
10            steps=data_handler.inferred_steps)  
11  
12      ## I am hook  
13      callbacks.on_train_begin()  
14      training_logs = None  
15      # Handle fault-tolerance for multi-worker.  
16      # TODO(omalleyt): Fix the ordering issues that mean this has to  
17      # happen after `callbacks.on_train_begin`.  
18      data_handler._initial_epoch = (  # pylint: disable=protected-access  
19          self._maybe_load_initial_epoch_from_ckpt(initial_epoch))  
20      for epoch, iterator in data_handler.enumerate_epochs():  
21        self.reset_metrics()  
22        callbacks.on_epoch_begin(epoch)  
23        with data_handler.catch_stop_iteration():  
24          for step in data_handler.steps():  
25            with trace.Trace(  
26                'TraceContext',  
27                graph_type='train',  
28                epoch_num=epoch,  
29                step_num=step,  
30                batch_size=batch_size):  
31              ## I am hook  
32              callbacks.on_train_batch_begin(step)  
33              tmp_logs = train_function(iterator)  
34              if data_handler.should_sync:  
35                context.async_wait()  
36              logs = tmp_logs  # No error, now safe to assign to logs.  
37              end_step = step + data_handler.step_increment  
38              callbacks.on_train_batch_end(end_step, logs)  
39        epoch_logs = copy.copy(logs)  
40  
41        # Run validation.  
42  
43        ## I am hook  
44        callbacks.on_epoch_end(epoch, epoch_logs)  
45

3.2 mmdetection

mmdetection是一个目标检测的开源框架，集成了许多不同的目标检测深度学习算法（pytorch版），如faster-rcnn, fpn, retianet等。里面也大量使用了hook，暴露给应用实现流程中具体部分。

详见https://github.com/open-mmlab/mmdetection

这里看一个训练的调用例子(摘录)（https://github.com/open-mmlab/mmdetection/blob/5d592154cca589c5113e8aadc8798bbc73630d98/mmdet/apis/train.py）

1def train_detector(model,  
2                   dataset,  
3                   cfg,  
4                   distributed=False,  
5                   validate=False,  
6                   timestamp=None,  
7                   meta=None):  
8    logger = get_root_logger(cfg.log_level)  
9  
10    # prepare data loaders  
11  
12    # put model on gpus  
13  
14    # build runner  
15    optimizer = build_optimizer(model, cfg.optimizer)  
16    runner = EpochBasedRunner(  
17        model,  
18        optimizer=optimizer,  
19        work_dir=cfg.work_dir,  
20        logger=logger,  
21        meta=meta)  
22    # an ugly workaround to make .log and .log.json filenames the same  
23    runner.timestamp = timestamp  
24  
25    # fp16 setting  
26    # register hooks  
27    runner.register_training_hooks(cfg.lr_config, optimizer_config,  
28                                   cfg.checkpoint_config, cfg.log_config,  
29                                   cfg.get('momentum_config', None))  
30    if distributed:  
31        runner.register_hook(DistSamplerSeedHook())  
32  
33    # register eval hooks  
34    if validate:  
35        # Support batch_size > 1 in validation  
36        eval_cfg = cfg.get('evaluation', {})  
37        eval_hook = DistEvalHook if distributed else EvalHook  
38        runner.register_hook(eval_hook(val_dataloader, **eval_cfg))  
39  
40    # user-defined hooks  
41    if cfg.get('custom_hooks', None):  
42        custom_hooks = cfg.custom_hooks  
43        assert isinstance(custom_hooks, list), \  
44            f'custom_hooks expect list type, but got {type(custom_hooks)}'  
45        for hook_cfg in cfg.custom_hooks:  
46            assert isinstance(hook_cfg, dict), \  
47                'Each item in custom_hooks expects dict type, but got ' \  
48                f'{type(hook_cfg)}'  
49            hook_cfg = hook_cfg.copy()  
50            priority = hook_cfg.pop('priority', 'NORMAL')  
51            hook = build_from_cfg(hook_cfg, HOOKS)  
52            runner.register_hook(hook, priority=priority)  
53

4. 总结

本文介绍了hook的概念和应用，并给出了python的实现细则。希望对比有帮助。总结如下：

hook函数是流程中预定义好的一个步骤，没有实现
挂载或者注册时，流程执行就会执行这个钩子函数
回调函数和hook函数功能上是一致的
hook设计方式带来灵活性，如果流程中有一个步骤，你想让调用方来实现，你可以用hook函数

作者简介：wedo实验君, 数据分析师；热爱生活，热爱写作

-------------------********************************** End **********-------------**-----********-**********************************

往期精彩文章推荐：

欢迎各位大佬点击链接加入群聊【helloworld开发者社区】：https://jq.qq.com/?_wv=1027&k=mBlk6nzX进群交流IT技术热点。

本文转自 https://mp.weixin.qq.com/s/1ELA1lu5x19OFOdjB_iBJQ，如有侵权，请联系删除。