在上一篇文章,我们讲述了该框架的Trainer结构,其实该框架中最好玩的应该是Hook的结构,接下来我们详细来讲Hook
可以看到在Trainer类中,会存储一个hook的list,list中是不同功能的hook,然后在训练的不同阶段,会挨个调用所有hook,如果该hook在这个阶段需要工作,则进行该hook的运行。那么hook是怎么获取trainer过程中的数据结果之类的呢
hooks = [h for h in hooks if h is not None]
for h in hooks:
assert isinstance(h, HookBase)
# To avoid circular reference, hooks and trainer cannot own each other.
# This normally does not matter, but will cause memory leak if the
# involved objects contain __del__:
# See http://engineering.hearsaysocial.com/2013/06/16/circular-references-in-python/
h.trainer = weakref.proxy(self)
在TrainerBase中,我们可以看到这一段代码,该代码将所有的hook进行遍历之后,去掉空hook,然后将该Trainer对象的引用保存在hook.trainer中,这样hook就可以得到训练过程中结果
class HookBase:
"""
Base class for hooks that can be registered with :class:`TrainerBase`.
Each hook can implement 4 methods. The way they are called is demonstrated
in the following snippet:
.. code-block:: python
hook.before_train()
for iter in range(start_iter, max_iter):
hook.before_step()
trainer.run_step()
hook.after_step()
hook.after_train()
Notes:
1. In the hook method, users can access `self.trainer` to access more
properties about the context (e.g., current iteration).
2. A hook that does something in :meth:`before_step` can often be
implemented equivalently in :meth:`after_step`.
If the hook takes non-trivial time, it is strongly recommended to
implement the hook in :meth:`after_step` instead of :meth:`before_step`.
The convention is that :meth:`before_step` should only take negligible time.
Following this convention will allow hooks that do care about the difference
between :meth:`before_step` and :meth:`after_step` (e.g., timer) to
function properly.
Attributes:
trainer: A weak reference to the trainer object. Set by the trainer when the hook is
registered.
"""
def before_train(self):
"""
Called before the first iteration.
"""
pass
def after_train(self):
"""
Called after the last iteration.
"""
pass
def before_step(self):
"""
Called before each iteration.
"""
pass
def after_step(self):
"""
Called after each iteration.
"""
pass
所有hook都是继承于HookBase这个类,HookBase抽象了各个step的之间的操作。
/detectron2/engine/hooks.py中定义了各种不同功能的hook,计算训练时间的hook,训练速率调整的hook,eval的hook等