【金猿技术展】模块注册机制——打造AI算法流水线
2021-12-23 20:45:07
  • 0
  • 0
  • 0

模块注册机制基于注册器和构建器,可以轻松地进行模块拓展,而不需要在原有框架代码中进行改动。使用注册器管理字符串到类的映射,其中支持注册的模块类型包括但不限于模型结构、数据预处理模块、优化器。使用构建器可以将配置文件转换成对应的模块,其提供了灵活的构建方式,使得可以构建定制化的训练流水线。模块注册机制的另一特性是采用分层策略划分模块的注册域,不仅能够避免跨项目间同名模块的冲突,也能够支持跨项目间模块的相互调用。

技术说明

模块注册和跨项目调用机制

开发算法工具库以解决人工智能基础理论知识与商用业务落地工程经验间脱节的问题。OpenMMLab基于在学术研究和行业落地的经验,对 AI算法进行标准化模块工具开发,大大降低具体行业业务人员的算法学习成本,将以工业化的人工智能研究颠覆传统手工作坊式的研究方式,提供先进的研发工具,使人工智能研究进入快车道。

算法核心组件作为开放算法体系的基础,为算法提供公共的架构支持和模块实现。核心组件主要分为三个部分:基础支持库、通用算子和训练流程框架。基础支持库提供了友好的文件读写接口、图像视频处理接口、基础工具库等,通用算子提供了常用算子的高效的 GPU 实现,训练流程框架提供了通用的算法训练流程支持,使得体系内的算法能够在相同的架构中进行训练。

在 OpenMMLab 开源算法体系的整体架构中,突出两个特点:模块化和简洁化。模块化让用户可以更灵活地进行排列组合,以及开发新的方法和模块;简洁化让用户不用过多关注和核心逻辑无关的事情,将辅助功能交由框架自行处理。

OpenMMLab 开源算法体系通过配置文件、注册器和构建器实现模块化的基础支持,然后在模块层面将数据、模型等进行拆解和抽象,进行模块化的实现,在最顶层通过执行器来实现流程的简洁化。

注册器提供了在不修改算法框架核心代码的前提下进行模块拓展的能力。通过管理字符串到类的映射,可以实现从配置文件直接构建模块。在开放算法体系中,不论是模型结构,数据预处理组件,还是优化器,甚至顶层的训练流程,都是通过这种方式来构建的,给用户提供了极大的拓展空间。

图2-2 MMCV 中注册器的构建机制

基于注册器和构建器,可以轻松地进行模块拓展,而不需要在原有框架代码中进行改动,这也是OpenMMLab 推荐的进行算法开发的流程。

在 OpenMMLab 中,Registry 类可以提供一种完全相似的对外装饰函数来管理构建不同的组件,例如 backbones、head 和 necks 等等,Registry 类内部维护的是一个全局 key-value 对。通过 Registry 类,用户可以通过字符串方式实例化任何想要的模块。例如在 Faster R-CNN 的 backbone 模块实例化时,可以采用如下配置:

backbone=dict(

type='ResNet', # 待实例化的类名

depth=50, # 后面的都是对于的类初始化参数

num_stages=4,

out_indices=(0, 1, 2, 3),

frozen_stages=1,

norm_cfg=dict(type='BN', requires_grad=True),

norm_eval=True,

style='pytorch'),

Registry 类最大好处是:解耦性强、可扩展性强,代码更易理解。

回到 Registry 类本身,有如下几种用法:

# 0. 先构建一个全局的 CATS 注册器类

CATS = mmcv.Registry('cat')

# 通过装饰器方式作用在想要加入注册器的具体类中

# 1. 不需要传入任何参数,此时默认实例化的配置字符串是 str (类名)

@CATS.register_module()

class BritishShorthair:

pass

# 类实例化

CATS.get('BritishShorthair')(**args)

# 2.传入指定 str,实例化时候只需要传入对应相同 str 即可

@CATS.register_module(name='Siamese')

class SiameseCat:

pass

# 类实例化

CATS.get('Siamese')(**args)

# 3.如果出现同名 Registry Key,可以选择报错或者强制覆盖

# 如果指定了 force=True,那么不会报错

# 此时 Registry 的 Key 中,Siamese2Cat 类会覆盖 SiameseCat 类

# 否则会报错

@CATS.register_module(name='Siamese',force=True)

class Siamese2Cat:

pass

# 类实例化

CATS.get('Siamese')(**args)

# 4. 可以直接注册类

class Munchkin:

pass

CATS.register_module(Munchkin)

# 类实例化

CATS.get('Munchkin')(**args)

目前 Registry 类一共有4种用法,方便不同场景下注册。

跨库调用

OpenMMLab 使用一系列算法库支持了视觉领域绝大部分任务。这些任务之间在模块组件上存在着一定的共通之处,比如在检测、分割等任务中,网络包含一个用于提取图片特征的主干网络,该主干网络的权重在不同任务间有通用性。因此在进行检测、分割等任务之前,可以首先使用主干网络进行分类任务的训练。利用分类任务庞大的 ImageNet 数据集,可以提高主干网络的特征提取能力,并提高模型训练的收敛速度。

利用 Registry 机制,OpenMMLab 的算法库能够便捷地完成如上所述的,跨算法库的训练流程,使各个算法库开发的模块有机结合在一起。以检测任务的主干网络跨库调用为例,在 MMDetection 中,检测算法 Yolo V3 的一种默认配置是使用 MobileNet V2 主干网络进行图像特征提取。在实验中,如有需求使用其他主干网络进行特征提取,如 MobileNet V3,而 MMDetection 算法库中没有该主干网络的实现,我们可以利用跨库调用,直接从实现了该主干网络的算法库中调用对应的模块,而不必重新实现需要的网络。在这里,MobileNet V3 主干网络在分类代码库 MMClassification 中有对应的实现,使用如下配置即可完成模型的跨库调用:

# 直接继承 yolo v3 的原始配置

_base_ = "./yolov3_mobilenetv2_320_300e_coco.py"

# 因为 MMDetection 中没有导入 MMClassification 算法库

# 为了完成相关的模块注册,需要调用 custom_imports 来指定额外的导入

custom_imports=dict(imports='mmcls.models', allow_failed_imports=False)

model = dict(

backbone=dict(

# 使用 "scope.type" 的语法,指定从 MMClassification 中寻找需要的模块

type='mmcls.MobileNetV3',

# MobileNet V3 的其他设置

arch='large',

out_indices=(5, 11, 14),

init_cfg=dict(

type='Pretrained',

checkpoint='mmcls://mobilenet_v3_large'),

# 配置文件与继承的配置文件中相同字段的字典,默认会融合

# 这里使用 `_delete_` 来删除继承的配置文件中的其他配置

_delete_=True),

# 主干网络发生变化,其他相应的配置也需要改变

neck=dict(in_channels=[160, 112, 40])

跨代码库调用机制

自定义模块导入

利用注册器机制,在导入相对应的模块时,模块中所包含类自动地被注册到 Registry 中的。而具体导入模块的操作,被放在了代码执行的入口程序中,一般为 tools/train.py 和 tools/test.py,分别对应了训练和测试。

以 MMClassification 中的 tools/train.py 为例:

...

from mmcv.runner import get_dist_info, init_dist

...

from mmcls.datasets import build_dataset

from mmcls.models import build_classifier

• 通过导入 mmcv.runner 包,完成了 mmcv/runner/__init__.py 中一系列执行器、钩子、优化器等类的注册。

• 通过导入 mmcls.datasets 包,完成了mmcls/datasets/__init__.py 中一系列数据集的注册。

• 通过导入mmcls.models 包,完成了mmcls/models/__init__.py中一系列主干网络、颈部头部函数的注册。

模块的的导入和类的注册不一定仅发生在入口程序的最外层,比如与数据处理和增强相关的注册器 PIPELINES 是入口程序在执行 build_dataset 时,在 mmcls/datasets/base_dataset.py 中进行相关类的注册。

通常,如果不存在明显的依赖,一个代码库中的入口程序不会导入其他代码库的模块,因而其他代码库中的类也不会被注册。为了实现跨库调用,MMCV 在解析配置文件时,提供了显式导入自定义模块的 custom_imports 接口。

custom_imports=dict(imports='mmcls.models', allow_failed_imports=False)

通过在配置文件中加入上述代码,MMCV 会在解析时借助 Python 内置的 importlib 包中的import_module 函数实现其他代码库模块的导入,进而完成对应的一系列类的注册。

Regsitry 中的 scope 机制

在进行跨库调用时,需要使用 scope.type 的语法显示指明要调用的类所在的 scope,例如在type='mmcls.MobileNetV3' 中,scope 为 “mmcls”。

通过引入 scope 机制的引入,Registry 能够更灵活地处理不同代码库之间模块的命名冲突,例如在MMDetection 和 MMClassification 中均注册了 MobileNetV2 主干网络,而 Registry 机制不允许同一个注册器中注册两个相同名字的类。因此,每一个代码库在独立开发,需要注册自己的模型时,都会注册到代码库自己的 Regsitry 中,而不是注册到 MMCV 统一的 Registry 中,通过这种方式来避免与其他代码库产生冲突。

与此同时,每个代码库同类模块的 Registry 并非独立的,而是以 MMCV 中某个统一的 Registry 为父 Regsitry,从而形成如图所示的树状结构。

由于 MMCV 能够获得所有挂靠在主 Registry 下的各个子 Registry 的 scope,因而允许用户显式地指定 scope,并定位对应代码库的 Registry,从而实现跨库调用。

开发团队

●带队负责人:陈恺

陈恺博士,OpenMMLab负责人。上海人工智能实验室青年科学家,陈博士本科毕业于清华大学,在香港中文大学获得博士学位,在计算机视觉顶级会议和期刊上发表了十余篇论文,谷歌学术引用超过 2000,并多次在国际学术竞赛中获奖。所负责的OpenMMLab 开源项目在 GitHub 上获得了超过 40000 star,是深度学习时代极具影响力的计算机视觉开源算法体系,在学术界和工业界产生了广泛影响。

●其他重要成员:吕成器、张文蔚、周再达

●隶属机构:上海人工智能实验室&商汤科技

上海人工智能实验室是我国人工智能领域的新型科研机构,开展战略性、原创性、前瞻性的科学研究与技术攻关,突破人工智能的重要基础理论和关键核心技术,打造“突破型、引领型、平台型”一体化的大型综合性研究基地,支撑我国人工智能产业实现跨越式发展,目标建成国际一流的人工智能实验室,成为享誉全球的人工智能原创理论和技术的策源地。

作为人工智能软件公司,商汤科技以“坚持原创,让AI引领人类进步”为使命,“以人工智能实现物理世界和数字世界的连接,促进社会生产力可持续发展,并为人们带来更好的虚实结合生活体验”为愿景,旨在持续引领人工智能前沿研究,持续打造更具拓展性更普惠的人工智能软件平台,推动经济、社会和人类的发展,并持续吸引及培养顶尖人才,共同塑造未来。

相关评价

OpenMMLab系列开源框架对新手来说太好了,自己编程也会按照这些方法来,潜移默化,写出来的代码赏心悦目。

——华中科技大学用户free

研一AI小白一枚,喵喵lab让我少走弯路,API很好用,文档清晰,代码结构也对理解其他开源代码很有用。

——西南交通大学用户Jin

OpenMMLab非常方便,很多方向的相关算法统一易用,方便我们测试修改,快速跑通。

——中科院宁波材料所用户Fan

 
最新文章
相关阅读