用 OpenMMLab 轻松搭建主干网络，多种视觉任务一网打尽

OpenMMLab 框架几乎全面覆盖了深度学习视觉任务的方方面面。针对每一个具体的任务，我们都提供了一个相应的算法库，如用于分类任务的 MMClassification，用于检测任 MMDetection 和用于分割任务的 MMSegmentation 等等。

随之而来的一个问题是，如果我希望将研究的方法应用于多个任务，在多个任务上进行评测，是不是要把各个算法库都 fork 一份？更进一步地，如果我希望将我依赖 OpenMMLab 框架的代码开源出去，应该怎么做呢？

这里，以我们最为常见的跨任务研究——搭建主干网络为例，介绍一下如何快速开发一个新的主干网络，并分别使用 mmcls、mmdet 和 mmseg 轻松实现分类、检测和分割任务上的评测。

为了方便大家轻松上手，我们提供了一个示例仓库：github.com/mzr1996/bac…

在示例仓库中，我们以 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 一直致力于为大家提供面向视觉深度学习各种任务的工具箱。与此同时，面向各种任务的算法库又是有机统一的，本文介绍了一种联合使用多个算法库的实践方式，欢迎大家下载试用我们的示例仓库。另外，我们也会在github.com/open-mmlab/… 中持续更新各种使用 MIM 完成各种任务的示例，欢迎大家关注~