如何基于 MMCV 走上开源大佬之路

2022-01-18 09:07:31 浏览数 (1)

MMCV是一款计算机视觉基础库,为 OpenMMLab 的十几个算法库支持的超过15个训练任务提供了统一的抽象训练 API 、常用的神经网络模块封装、和基础的图像/视频处理算子等模块。

一直以来都听到社区的朋友们问如何基于 MMCV 或者 OpenMMLab 项目搭建自己的项目,或者经常有听到朋友们问某某某代码或者文件去哪里找

在今天这篇文章里,我们就来介绍介绍如何基于 MMCV 开发一个规模较大、全新方向的算法库。

同时,本文也可以被认为是之前 OpenMMLab 各个算法库中共性的归纳和总结, 也就是说,各个算法库都基本遵循了下文所描述的内容,理解了它,你也可以更加轻松地理解并使用任意一个 OpenMMLab 的算法库啦!

二话不说,直接上干货!

本文内容

代码组织结构

算法库入口 tools

功能性代码 packages

模型定义文件 configs

单元测试 tests

文档说明 docs

其他配置项

算法库设计小技巧

Dataset 的输入输出流

模型的输入输出流

1. 代码组织结构

首先是代码的组织架构,OpenMMLab 各个算法库的文件结构大致如下图所示,包含核心库 package 和工具性代码的目录 tools、文档、CI/CD 和开发相关的配置文件。

算法库入口 tools

当你在 clone 了当前算法库后,一般使用 tools 下的脚本来进行模型训练、测试、推理和部署等工作。

tools 文件夹提供了整个算法库训练、测试和推理的的脚本,并提供了一些必要的工具性脚本,例如在 MMDet 中有如下脚本:

- deployment 文件夹中放置模型部署的脚本

- analysis_tools 放置 log/模型的分析脚本

- data_converters 放置数据集预处理脚本

- train.py/test.py, slurm_train.sh/dist_train.sh:

一些常见的命令,模型的训练测试入口

- misc 放置一些杂项

功能性代码 packages

OpenMMLab 的算法库都支持 pip install,在安装过程中,package 文件夹下的代码会被安装进环境里,供下游任务 import。

安装完成后,拓展项目就可以采用如 import mmedit 或者 import mmdet 来使用 MMEditing 和 MMDetection 中的各类型模块。

和已有的 OpenMMLab 算法库一样,新的算法库也应当采用模块化设计(modular design)并使用注册器来管理具有相同功能的模块,如使用 BACKBONE registry 来管理 ResNet,ResNext 和 SE-ResNet 等主干网络。同时,相同功能的模块会放置在同一文件夹内方便查找。

OpenMMLab 算法库中的 package 部分都使用了类似的代码架构。他们的内容与作用如下:

apis:提供了训练、测试和推理的高层次 API,被 tools 中的启动脚本调用;

core:实现了算法库中的核心模块,包括基础的数据结构,evaluation 相关的实现,和自定义的一些 hooks 等;

models: 实现了算法库中的各类算法和算法依赖的loss、head 等神经网络模块,算法 models 部分根据不同任务中的典型结构有所区别,如 MMGeneration 中直接采用了 architectures,MMDetection 中有 detectors,MMSegmentation 中有 segmentors,这些文件夹中一般包含的是完成某个任务的完整具体模型;

datasets: 提供了数据集和 dataloader 的相关实现,其中算法训练过程中的 data augmentation 以及测试时需要对图片进行的前处理都放置在 data_pipelines 文件夹中;

utils:提供了一些通用的工具性函数,如 get_root_logger 等。

模型定义文件 configs

Configs 文件夹定义了代码库中实现的各类算法模型,在使用 tools 下的脚本进行模型训练和测试时,模型训练过程中的各种配置都是通过 configs 的文件来设置的。

单元测试 tests

除了功能性代码,算法库中还应当包含单元测试代码,单元测试的目录结构应该和 package 的目录及文件结构基本保持一致,尽量做到 package 中的一个文件/模块在 tests目录下都能找到一一对应的测试文件与测试模块。

文档说明 docs

Docs 文件夹中包含了算法库整体文档的入口,我们一般采用 sphinx 来配置文档,并且文档中一般会包含基本的 tutorial 来教用户如何安装、使用算法库。常见的必备文档还包括:

- changelog.md 来记录算法库每次新版本的变更;

- compatibility.md 记录算法库每次版本变更时发生的 BC-breaking;

- 项目中整体的 model zoo 介绍;

具体内容可参考 MMDetection 中的 changelog 、 compatibility 和 model zoo 文档。

其他配置项

其他配置项目包括:

- .pre-commit-config.yaml 定义了代码在提交 commit 时要经过的静态检查;

- setup.cfg 定义的代码静态检查时使用的规则;

- .github目录定义了 github 提交的 workflow 和 Issue/PR 模板。

2. 算法库设计小技巧

新的算法库开始设计时,应当着重定义两大模块 dataset 和 model 的抽象接口,如果是基于已有算法库的拓展,一般项目都会保持和依赖算法库一样的输入输出流并尽可能复用已有模块,避免版本不兼容的风险。

在新的算法库设计中,OpenMMLab 的算法库一般主要考虑两个大模块之间的抽象接口:数据和模型。

Dataset 的输入输出流

一般来说,OpenMMLab 算法库里的数据集在训练过程中执行的抽象过程如下图所示,在完成初始化后,dataset 会加载数据集的 meta 文件,并将 meta 文件中的各类信息重新组织成 data pipelines 的输入,经过 data pipelines 对数据进行各类型处理后,将数据输出给模型。

因此,在设计的时候需要:

- 确定 data pipelines 的输入:将各个数据集的数据结构组织成一个统一的内部中间格式,并传给 data pipelines;

- 确定各个 Pipeline 中流通的数据字段,以 MMDetection 中的 data pipeline 为例子,如下面代码所示,Data pipeline 会根据数据集的 annotation 文件,加载对应的图片和图片标注,并对图片做一系列 augmentation 操作。

以 MMDetection 为例

https://mmdetection.readthedocs.io/en/latest/tutorials/data_pipeline.html,每个 pipeline 都会给增加或者调整某些字段,在组合pipeline的时候应当小心(想起来目前的 pipeline 还没有做输入的字段检查容易踩坑,要给帅气的开发小哥哥加需求!)

代码语言:javascript复制
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
    dict(type='RandomFlip', flip_ratio=0.5),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size_divisor=32),
    dict(type='MyTransform'),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),
]

- 确定 data pipelines 的最终输入,需要抽象出用于模型训练的数据字段的最小集合,如 gt_bboxes, gt_labels, gt_masks,gt_seg,并且设计 Collect 和DefaultFormatBundle 模块在 data pipelines 的末尾收集必要字段整合成数据流的最终输出。这个输出会给到算法模型进行训练和推理。

模型的输入输出流

数据集 data pipelines 的输出最终会给到模型,模型测试的输出最终也会给到 dataset 的 evaluate 函数。

因此,这部分接口也需要提前进行定义,具体包括:模型输出的格式、对模型输出结果处理后的格式,以及evaluate 输入的接口。

模型的输出结果会经历的流程如下:

数据集会首先接受模型的输出,并对模型的输出进行 reformat,reformat 完成后就可以对结果进行评测,或者将结果 dump 下来,用于打比赛时测试结果的提交

本文简单介绍了基于 MMCV 开发一个全新算法库时代码的组织结构以及要考虑到的设计。

因为这同时也是 OpenMMLab 已有算法库中常见的文件结构和抽象逻辑,也希望本文能加深小伙伴们对 OpenMMLab 算法库的理解,希望以后再也不用担心找不到某个代码实现啦~

欢迎大家 Star :https://github.com/open-mmlab/mmcv

0 人点赞