OpenMMLab 框架几乎全面覆盖了深度学习视觉任务的方方面面。针对每一个具体的任务,我们都提供了一个相应的算法库,如用于分类任务的 MMClassification,用于检测任 MMDetection 和用于分割任务的 MMSegmentation 等等。
随之而来的一个问题是,如果我希望将研究的方法应用于多个任务,在多个任务上进行评测,是不是要把各个算法库都 fork 一份?更进一步地,如果我希望将我依赖 OpenMMLab 框架的代码开源出去,应该怎么做呢?
这里,以我们最为常见的跨任务研究——搭建主干网络为例,介绍一下如何快速开发一个新的主干网络,并分别使用 mmcls、mmdet 和 mmseg 轻松实现分类、检测和分割任务上的评测。
为了方便大家轻松上手,我们提供了一个示例仓库:https://github.com/mzr1996/backbone-example
在示例仓库中,我们以 ConvNeXt 为例,展示了搭建主干网络并进行基准测试所需要的文件。那么下面,让我们以这个示例仓库为例,一步一步了解如何搭建自己的主干网络吧。
配置环境
首先,我们需要配置 Python 和 PyTorch 环境,相信对 PyTorch 驾轻就熟的大家一定已经准备好了这些环境。需要注意的是,我们要求 Python >= 3.6,PyTorch >= 1.5.0,更新的版本当然也没有问题。
除此之外,我们还需要以下 OpenMMLab 仓库:
- MIM: 一个用于管理 OpenMMLab 包和实验的命令行工具
- MMCV: OpenMMLab 计算机视觉基础库
- MMClassification: OpenMMLab 图像分类工具箱和基准测试。除了分类任务,它同时用于提供多样的主干网络
- MMDetection: OpenMMLab 检测工具箱和基准测试
- MMSegmentation: OpenMMLab 语义分割工具箱和基准测试
假设你已经准备好了 Python 和 PyTorch 环境,那么只需要下面两行命令,就可以配置好软件环境。
pip install openmim mmcls mmdet mmsegmentation
mim install mmcv-full
数据准备
假设我们希望在 ImageNet 上进行分类任务基准测试,在 ADE20K 上进行语义分割任务基准测试,在 COCO 上进行检测任务基准测试,那么就需要按照以下格式组织好数据文件。
data/
├── imagenet
│ ├── train
│ ├── val
│ └── meta
│ ├── train.txt
│ └── val.txt
├── ade
│ └── ADEChallengeData2016
│ ├── annotations
│ └── images
└── coco
├── annotations
│ ├── instance_train2017.json
│ └── instance_val2017.json
├── train2017
└── val2017
这里,我们只列举了用于训练和验证的必要文件。如果你希望在更多数据集或任务上进行基准测试,比如使用 MMDetection 进行全景分割,只需要按照 MMDetection 的需要组织对应的数据集即可。另外,MMSegmentation 还提供了一篇教程讲解如何组织语义分割相关数据集。
实现你的主干网络
准备工作一切就绪,接下来让我们开始一步步实现和训练我们的主干网络吧~
首先是实现:
1.在 models
文件夹下创建你的主干网络文件。例如示例仓库中的 models/convnext.py
。在这个文件中,只需要使用 PyTorch 实现你的主干网络即可,额外需要注意的只有以下两点:
- 主干网络和主要模块需要继承
mmcv.runner.BaseModule
。这个BaseModule
是torch.nn.Module
的子类, 其行为几乎完全一致,除了它额外支持使用init_cfg
参数指定包括预训练模型在内的初始化方法。 - 使用如下的装饰器结构,将你的主干网络类注册至
mmcls.models.BACKBONES
注册器。
from mmcv.runner import BaseModule
from mmcls.models import BACKBONES
@BACKBONES.register_module(force=True)
class MyBackbone(BaseModule):
...
注册器是什么?看看 这个!
2.[可选] 如果你希望为某些特定的任务增加一些额外的模块、组件,可以参照 models/det/layer_decay_optimizer_constructor.py
添加至对应的文件夹中。
3.将你添加的主干网络、自定义组件添加至 models/__init__.py
文件中。
添加配置文件
如果你还不太清楚配置文件是怎么回事,这篇教程应该可以帮到你。
通常来说,我们使用若干基础配置文件,通过继承的方式来组织一个实验的配置文件。这些基础配置文件一般包括模型、数据集、优化策略和运行配置。你也可以在配置文件中对基础配置文件进行覆盖,或者不使用基础配置文件,全部写到一个配置文件里。
在本示例库中,我们提供了一套比较常用的基础配置文件,你可以在从以下位置访问链接找到更多有用的基础配置文件:MMClassification,MMDetection, MMSegmentation。
对每个不同的任务,对应的配置文件可以放在放在 configs
目录中不同的子文件夹。需要注意的是,在配置文件的 model
部分,为了能够使 MMDetection 和 MMSegmentation 能够调用注册在 mmcls.models.BACKBONES
的主干网络,我们需要在 det 和 seg 的配置文件中额外指定主干网络注册器的位置,这只需要在 type
中添加 mmcls.
前缀即可,如:
model = dict(
backbone=dict(type='mmcls.MyBackbone', ...),
...
)
想要了解跨仓库调用模块的更多内容,可以参见《MMDet居然能用MMCls的Backbone?论配置文件的打开方式》
训练和测试
依托于 OpenMMLab 提供的统一实验管理工具—— MIM,我们在实现了模型、编写了配置文件之后,不需要写任何 Python 脚本就可以进行不同任务的训练和测试。
首先,我们需要将当前仓库的目录添加到 PYTHONPATH
环境变量中,这样 Python 才可以找到我们的模型文件,在仓库根目录运行如下指令:
export PYTHONPATH=`pwd`:$PYTHONPATH
单机单 GPU
# 训练分类模型
mim train mmcls $CONFIG --work-dir $WORK_DIR
# 测试分类模型
mim test mmcls $CONFIG -C $CHECKPOINT --metrics accuracy --metric-options "topk=(1, 5)"
# 训练目标检测/实例分割模型
mim train mmdet $CONFIG --work-dir $WORK_DIR
# 测试目标检测/实例分割模型
mim test mmdet $CONFIG -C $CHECKPOINT --eval bbox segm
# 训练语义分割模型
mim train mmseg $CONFIG --work-dir $WORK_DIR
# 测试语义分割模型
mim test mmseg $CONFIG -C $CHECKPOINT --eval mIoU
一些参数的说明:
$CONFIG
:configs/
文件夹中的配置文件。$WORK_DIR
:用于保存日志和权重文件的文件夹$CHECKPOINT
:权重文件路径
单机多 GPU (以 4 GPU 为例)
# 训练分类模型
mim train mmcls $CONFIG --work-dir $WORK_DIR --launcher pytorch --gpus 4
# 测试分类模型
mim test mmcls $CONFIG -C $CHECKPOINT --metrics accuracy --metric-options "topk=(1, 5)" --launcher pytorch --gpus 4
# 训练目标检测/实例分割模型
mim train mmdet $CONFIG --work-dir $WORK_DIR --launcher pytorch --gpus 4
# 测试目标检测/实例分割模型
mim test mmdet $CONFIG -C $CHECKPOINT --eval bbox segm --launcher pytorch --gpus 4
# 训练语义分割模型
mim train mmseg $CONFIG --work-dir $WORK_DIR --launcher pytorch --gpus 4
# 测试语义分割模型
mim test mmseg $CONFIG -C $CHECKPOINT --eval mIoU --launcher pytorch --gpus 4
结语
OpenMMLab 一直致力于为大家提供面向视觉深度学习各种任务的工具箱。与此同时,面向各种任务的算法库又是有机统一的,本文介绍了一种联合使用多个算法库的实践方式,欢迎大家下载试用我们的示例仓库。另外,我们也会在https://github.com/open-mmlab/mim-example 中持续更新各种使用 MIM 完成各种任务的示例,欢迎大家关注~