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

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

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

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

为了方便大家轻松上手

我们提供了一个示例仓库

https://github.com/mzr1996/backbone-example

（文章末尾阅读原文可直达）

在示例仓库中，我们以 ConvNeXt 为例，展示了搭建主干网络并进行基准测试所需要的文件。那么下面，让我们以这个示例仓库为例，一步一步了解如何搭建自己的主干网络吧。

配置环境

首先，我们需要配置 Python 和 PyTorch 环境，相信对 PyTorch 驾轻就熟的大家一定已经准备好了这些环境。需要注意的是，我们要求 Python >= 3.6，PyTorch >= 1.5.0，更新的版本当然也没有问题。

PyTorch 网址：

https://pytorch.org/get-started/locally/

除此之外，我们还需要以下 OpenMMLab 仓库：

MIM：

https://github.com/open-mmlab/mim

一个用于管理 OpenMMLab 包和实验的命令行工具。

MMCV：

https://github.com/open-mmlab/mmcv

OpenMMLab 计算机视觉基础库。

MMClassification：

https://github.com/open-mmlab/mmclassification

OpenMMLab 图像分类工具箱和基准测试。除了分类任务，它同时用于提供多样的主干网络。

MMDetection：

https://github.com/open-mmlab/mmdetection

OpenMMLab 检测工具箱和基准测试。

MMSegmentation：

https://github.com/open-mmlab/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 还提供了一篇教程讲解如何组织语义分割相关数据集。

教程网址：

https://mmsegmentation.readthedocs.io/zh_CN/latest/dataset_prepare.html

实现你的主干网络

准备工作一切就绪，接下来让我们开始一步步实现和训练我们的主干网络吧~

首先是实现：

1）在 models 文件夹下创建你的主干网络文件。例如示例仓库中的 models/convnext.py。在这个文件中，只需要使用 PyTorch 实现你的主干网络即可，额外需要注意的只有以下两点：

a. 主干网络和主要模块需要继承 mmcv.runner.BaseModule。这个 BaseModule 是 torch.nn.Module 的子类， 其行为几乎完全一致，除了它额外支持使用 init_cfg 参数指定包括预训练模型在内的初始化方法。

b. 使用如下的装饰器结构，将你的主干网络类注册至 mmcls.models.BACKBONES 注册器。

models/convnext.py 网址：

https://github.com/mzr1996/backbone-example/blob/master/models/convnext.py

注册器是什么？看这里：

https://mmcv.readthedocs.io/zh_CN/latest/understand_mmcv/registry.html

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 文件中。

models/det/layer_decay_optimizer_constructor.py 网址：

https://github.com/mzr1996/backbone-example/blob/master/models/det/layer_decay_optimizer_constructor.py

models/__init__.py 网址：

https://github.com/mzr1996/backbone-example/blob/master/models/__init__.py

添加配置文件

如果你还不太清楚配置文件是怎么回事，这篇教程应该可以帮到你。

教程网址：

https://mmclassification.readthedocs.io/zh_CN/latest/tutorials/config.html#config-file-structure

通常来说，我们使用若干基础配置文件，通过继承的方式来组织一个实验的配置文件。这些基础配置文件一般包括模型、数据集、优化策略和运行配置。你也可以在配置文件中对基础配置文件进行覆盖，或者不使用基础配置文件，全部写到一个配置文件里。

在本示例库中，我们提供了一套比较常用的基础配置文件，可以访问对应的算法库找到更多有用的基础配置文件。

对每个不同的任务，对应的配置文件可以放在放在 configs 目录中不同的子文件夹。需要注意的是，在配置文件的 model 部分，为了能够使 MMDetection 和 MMSegmentation 能够调用注册在 mmcls.models.BACKBONES 的主干网络，我们需要在 MMDetection 和 MMSegmentation 的配置文件中额外指定主干网络注册器的位置，这只需要在 type 中添加 mmcls. 前缀即可，如：

configs 网址：

https://github.com/mzr1996/backbone-example/blob/master/configs

model = dict(
    backbone=dict(type='mmcls.MyBackbone', ...),
    ...
)

想要了解跨仓库调用模块的更多内容，

训练和测试

依托于 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