详解目标检测(MMdetection)-HOOK机制

2020 年 9 月 22 日 极市平台
↑ 点击 蓝字  关注极市平台

作者丨却道@知乎
来源丨https://zhuanlan.zhihu.com/p/238130913
编辑丨极市平台

极市导读

 

本文作者分享了自己对MMdetetion中钩子(HOOK)机制的看法,包括了作用理解和典型案例等内容。>>加入极市CV技术交流群,走在计算机视觉的最前沿


最近做了一段时间的目标检测,不得不说检测这块还是相对比较复杂的,在熟悉项目的同时也确实学习到了很多有用的东西。MMdetetion是现在最著名、算法包最多并且使用人数最多的训练框架,其中的源码非常值得学习,今天总结下我对其中HOOK(钩子)机制的理解。

MMdetection最近更新很多,我以2.4.0版本的代码进行解读,分享自己的理解,也吸纳观众的点评。HOOK、Runer的定义在MMCV当中,MMdetection和MMCV是版本匹配的,我这里使用的是MMCV 1.1.2的代码。(HOOK相关的定义主要在MMCV中,下面用的代码都是摘自于MMCV)。

1.HOOK机制的作用

MMdetection中的HOOK可以理解为一种触发器,也可以理解为一种训练框架的架构规范,它规定了在算法训练过程中的种种操作,并且我们可以通过继承HOOK类,然后注册HOOK自定义我们想要的操作。

首先看一下HOOK的基类定义

# Copyright (c) Open-MMLab. All rights reserved.
from mmcv.utils import Registry

HOOKS = Registry('hook')


class Hook:

    def before_run(self, runner):
        pass

    def after_run(self, runner):
        pass

    def before_epoch(self, runner):
        pass

    def after_epoch(self, runner):
        pass

    def before_iter(self, runner):
        pass

    def after_iter(self, runner):
        pass

    def before_train_epoch(self, runner):
        self.before_epoch(runner)

    def before_val_epoch(self, runner):
        self.before_epoch(runner)

    def after_train_epoch(self, runner):
        self.after_epoch(runner)

    def after_val_epoch(self, runner):
        self.after_epoch(runner)

    def before_train_iter(self, runner):
        self.before_iter(runner)

    def before_val_iter(self, runner):
        self.before_iter(runner)

    def after_train_iter(self, runner):
        self.after_iter(runner)

    def after_val_iter(self, runner):
        self.after_iter(runner)

    def every_n_epochs(self, runner, n):
        return (runner.epoch + 1) % n == 0 if n > 0 else False

    def every_n_inner_iters(self, runner, n):
        return (runner.inner_iter + 1) % n == 0 if n > 0 else False

    def every_n_iters(self, runner, n):
        return (runner.iter + 1) % n == 0 if n > 0 else False

    def end_of_epoch(self, runner):
        return runner.inner_iter + 1 == len(runner.data_loader)

可以说基类函数中定义了许多我们在模型训练中需要用到的一些功能,如果想定义一些操作我们就可以继承这个类并定制化我们的功能,可以看到HOOK中每一个参数都是有runner作为参数传入的。关于Runner的作用下一篇文章接着说,简而言之,Runner是一个模型训练的工厂,在其中我们可以加载数据、训练、验证以及梯度backward等等全套流程。MMdetection在设计的时候也为runner传入丰富的参数,定义了一个非常好的训练范式。在你的每一个hook函数中,都可以对runner进行你想要的操作。

而HOOK是怎么嵌套进runner中的呢?其实是在Runner中定义了一个hook的list,list中的每一个元素就是一个实例化的HOOK对象。其中提供了两种注册hook的方法,register_hook是传入一个实例化的HOOK对象,并将它插入到一个列表中,register_hook_from_cfg是传入一个配置项,根据配置项来实例化HOOK对象并插入到列表中。当然第二种方法又是MMLab的开源生态中定义的一种基础方法mmcv.build_from_cfg了,无论在MMdetection还是其他MMLab开源的算法框架中,都遵循着MMCV的这套基于配置项实例化对象的方法。毕竟MMCV是提供了一个基础的功能,服务于各个算法框架,这也是为什么MMLab的代码高质量的原因。不仅仅是算法的复现,更是架构、编程范式的一种体现,真·代码如诗

def register_hook(self, hook, priority='NORMAL'):
        """Register a hook into the hook list.
        The hook will be inserted into a priority queue, with the specified
        priority (See :class:`Priority` for details of priorities).
        For hooks with the same priority, they will be triggered in the same
        order as they are registered.
        Args:
            hook (:obj:`Hook`): The hook to be registered.
            priority (int or str or :obj:`Priority`): Hook priority.
                Lower value means higher priority.
        """

        assert isinstance(hook, Hook)
        if hasattr(hook, 'priority'):
            raise ValueError('"priority" is a reserved attribute for hooks')
        priority = get_priority(priority)
        hook.priority = priority
        # insert the hook to a sorted list
        inserted = False
        # hook是分优先级插入到list中的,在MMdetection中不同的HOOK是有优先级的,为什么呢?稍后在hook的调用中解释哈
        for i in range(len(self._hooks) - 1-1-1):
            if priority >= self._hooks[i].priority:
                self._hooks.insert(i + 1, hook)
                inserted = True
                break
        if not inserted:
            self._hooks.insert(0, hook)

    def register_hook_from_cfg(self, hook_cfg):
        """Register a hook from its cfg.
        Args:
            hook_cfg (dict): Hook config. It should have at least keys 'type'
              and 'priority' indicating its type and priority.
        Notes:
            The specific hook class to register should not use 'type' and
            'priority' arguments during initialization.
        """

        hook_cfg = hook_cfg.copy()
        priority = hook_cfg.pop('priority''NORMAL')
        hook = mmcv.build_from_cfg(hook_cfg, HOOKS)
        self.register_hook(hook, priority=priority)

调用HOOK函数

def call_hook(self, fn_name):
        """Call all hooks.
        Args:
            fn_name (str): The function name in each hook to be called, such as
                "before_train_epoch".
        """

        for hook in self._hooks:
            getattr(hook, fn_name)(self)

可以看到HOOK是调用的时候是遍历List,然后根据HOOK的名字来调用。这也是为什么要区分优先级的原因,优先级越高的放在List的前面,这样就能更快地被调用。当你想用_before_run_epoch_来做A和B两件事情的时候,在runner里面就是调用一次self.before_run_epoch,但是先做A还是先做B,就是通过不同的HOOK的优先级来决定了。比如在evaluation的时候对需要做测试,但是测试前对参数做滑动平均。比如emaHOOK中的72行,也写明了要在测试之前做指数滑动平均。

def after_train_epoch(self, runner):
        """We load parameter values from ema backup to model before the
        EvalHook."""

        self._swap_ema_parameters()

checkpoint.py的HOOK中,同样也定义了after_train_epoch函数如下:

@master_only
    def after_train_epoch(self, runner):
        if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
            return

        runner.logger.info(f'Saving checkpoint at {runner.epoch + 1} epochs')
        if not self.out_dir:
            self.out_dir = runner.work_dir
        runner.save_checkpoint(
            self.out_dir, save_optimizer=self.save_optimizer, **self.args)

        # remove other checkpoints
        if self.max_keep_ckpts > 0:
            filename_tmpl = self.args.get('filename_tmpl''epoch_{}.pth')
            current_epoch = runner.epoch + 1
            for epoch in range(current_epoch - self.max_keep_ckpts, 0-1):
                ckpt_path = os.path.join(self.out_dir,
                                         filename_tmpl.format(epoch))
                if os.path.exists(ckpt_path):
                    os.remove(ckpt_path)
                else:
                    break

从测试代码中可以看到不同的HOOK虽然都是重写了after_train_epoch函数,但是调用的顺序还是先调用ema.py中的,然后再调用checkpoint.py中的after_train_epoch

resume_ema_hook = EMAHook(
        momentum=0.5, warm_up=0, resume_from=f'{work_dir}/epoch_1.pth')
    runner = _build_demo_runner()
    runner.model = demo_model
    # 设置了HIGHREST的优先级
    runner.register_hook(resume_ema_hook, priority='HIGHEST')
    checkpointhook = CheckpointHook(interval=1, by_epoch=True)
    runner.register_hook(checkpointhook)
    runner.run([loader, loader], [('train'1), ('val'1)], 2)

具体的优先级定义有以下7种,作为HOOK的类成员属性。具体定义在链接中。

+------------+------------+
| Level | Value |
+============+============+
| HIGHEST | 0 |
+------------+------------+
| VERY_HIGH | 10 |
+------------+------------+
| HIGH | 30 |
+------------+------------+
| NORMAL | 50 |
+------------+------------+
| LOW | 70 |
+------------+------------+
| VERY_LOW | 90 |
+------------+------------+
| LOWEST | 100 |
+------------+------------+

2.举一个简单的例子

最近打算好好锻炼身体,健康生活,努力工作,我打算让自己变得更加自律。我给自己定下了几个条例,每天吃早饭之前得晨练30分钟,运动完之后才会感觉充满活力。每天吃午饭之前我得跑上一个实验,吃完饭之后回来刚好可以看下中间结果,吃完午饭之后我感觉结果没问题我需要午休30分钟, 晚上下班前我如果没什么事再锻炼30分钟。秉承着这样的原则我给自己定义一个HOOK来规范我的生活。

  • 定义我的HOOK
import sys
class HOOK:

    def before_breakfast(self, runner):
        print('{}:吃早饭之前晨练30分钟'.format(sys._getframe().f_code.co_name))

    def after_breakfast(self, runner):
        print('{}:吃早饭之前晨练30分钟'.format(sys._getframe().f_code.co_name))

    def before_lunch(self, runner):
        print('{}:吃午饭之前跑上实验'.format(sys._getframe().f_code.co_name))

    def after_lunch(self, runner):
        print('{}:吃完午饭午休30分钟'.format(sys._getframe().f_code.co_name))

    def before_dinner(self, runner):
        print('{}: 没想好做什么'.format(sys._getframe().f_code.co_name))

    def after_dinner(self, runner):
        print('{}: 没想好做什么'.format(sys._getframe().f_code.co_name))

    def after_finish_work(self, runner, are_you_busy=False):
        if are_you_busy:
            print('{}:今天事贼多,还是加班吧'.format(sys._getframe().f_code.co_name))
        else:
            print('{}:今天没啥事,去锻炼30分钟'.format(sys._getframe().f_code.co_name))
  • 定义我的Runner
class Runner(object):
    def __init__(self, ):
        pass
        self._hooks = []

    def register_hook(self, hook):
        # 这里不做优先级判断,直接在头部插入HOOK
        self._hooks.insert(0, hook)

    def call_hook(self, hook_name):
        for hook in self._hooks:
            getattr(hook, hook_name)(self)

    def run(self):
        print('开始启动我的一天')
        self.call_hook('before_breakfast')
        self.call_hook('after_breakfast')
        self.call_hook('before_lunch')
        self.call_hook('after_lunch')
        self.call_hook('before_dinner')
        self.call_hook('after_dinner')
        self.call_hook('after_finish_work')
        print('~~睡觉~~')
  • 运行main函数,注册HOOK并且调用Runner.run()开启我的一天
from MyHook import HOOK
from MyRunner import Runner
runner = Runner()
hook = HOOK()
runner.register_hook(hook)
runner.run()
  • 得到的输出结果如下:
开始启动我的一天
before_breakfast:吃早饭之前晨练30分钟
after_breakfast:吃早饭之前晨练30分钟
before_lunch:吃午饭之前跑上实验
after_lunch:吃完午饭午休30分钟
before_dinner: 没想好做什么
after_dinner: 没想好做什么
after_finish_work:今天没啥事,去锻炼30分钟
~~睡觉~~

3.总结

MMdetection中的HOOK设计巧妙,很好地对算法训练、测试进行了抽象和解耦。每一个做上层算法模型的,都值得一看。感谢MMLab贡献这么优质的代码,让我等凡夫俗子醍醐灌顶。

除了HOOK之外,这个代码中还有很多优质的思想。比如Runner是怎么做到包办一切的?注册器这个中枢管理系统是怎么工作的?多卡训练的一些坑是怎么解决的?等等等等,我也在持续地学习和消化。路漫漫其修远兮,吾将上下而求索。

一个小题目:我的代码中每个函数输出的时候都会打印出这个函数名,这个可以用_装饰器_很方便地解决奥。装饰器这个东西在MMLab的系列项目中有大量的应用。其中对fp16的支持让大家赞不绝口。接下来有时间,对Runner、Register、装饰器这些东西好好盘一盘。


推荐阅读


来测一测你是几级的炼丹师?

添加极市小助手微信(ID : cvmart2),备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳),即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群:月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~

△长按添加极市小助手

△长按关注极市平台,获取 最新CV干货

觉得有用麻烦给个在看啦~   
登录查看更多
9

相关内容

最新《Transformers模型》教程,64页ppt
专知会员服务
306+阅读 · 2020年11月26日
专知会员服务
17+阅读 · 2020年9月6日
深度学习目标检测方法及其主流框架综述
专知会员服务
147+阅读 · 2020年6月26日
注意力图神经网络的多标签文本分类
专知会员服务
111+阅读 · 2020年3月28日
专知会员服务
109+阅读 · 2020年3月12日
MMDetection v2.0 训练自己的数据集
CVer
30+阅读 · 2020年8月9日
如何评价FAIR最新开源的Detectron2目标检测框架?
极市平台
18+阅读 · 2019年10月14日
我这两年的目标检测
极市平台
8+阅读 · 2019年9月18日
实战 | 源码入门之Faster RCNN
计算机视觉life
19+阅读 · 2019年4月16日
用PyTorch做物体检测和追踪
AI研习社
12+阅读 · 2019年1月6日
【干货】注意力机制(Attention)最新综述论文及相关源码
GAN生成式对抗网络
11+阅读 · 2018年11月16日
干货 | 基于深度学习的目标检测算法综述
AI科技评论
18+阅读 · 2018年9月1日
Arxiv
6+阅读 · 2019年4月4日
Arxiv
5+阅读 · 2018年4月17日
Arxiv
7+阅读 · 2018年3月19日
Arxiv
4+阅读 · 2018年3月19日
Arxiv
7+阅读 · 2018年1月24日
VIP会员
相关VIP内容
最新《Transformers模型》教程,64页ppt
专知会员服务
306+阅读 · 2020年11月26日
专知会员服务
17+阅读 · 2020年9月6日
深度学习目标检测方法及其主流框架综述
专知会员服务
147+阅读 · 2020年6月26日
注意力图神经网络的多标签文本分类
专知会员服务
111+阅读 · 2020年3月28日
专知会员服务
109+阅读 · 2020年3月12日
相关资讯
MMDetection v2.0 训练自己的数据集
CVer
30+阅读 · 2020年8月9日
如何评价FAIR最新开源的Detectron2目标检测框架?
极市平台
18+阅读 · 2019年10月14日
我这两年的目标检测
极市平台
8+阅读 · 2019年9月18日
实战 | 源码入门之Faster RCNN
计算机视觉life
19+阅读 · 2019年4月16日
用PyTorch做物体检测和追踪
AI研习社
12+阅读 · 2019年1月6日
【干货】注意力机制(Attention)最新综述论文及相关源码
GAN生成式对抗网络
11+阅读 · 2018年11月16日
干货 | 基于深度学习的目标检测算法综述
AI科技评论
18+阅读 · 2018年9月1日
相关论文
Arxiv
6+阅读 · 2019年4月4日
Arxiv
5+阅读 · 2018年4月17日
Arxiv
7+阅读 · 2018年3月19日
Arxiv
4+阅读 · 2018年3月19日
Arxiv
7+阅读 · 2018年1月24日
Top
微信扫码咨询专知VIP会员