MMPreTrain 算法库介绍

MMPretrain 是一个全新升级的预训练开源算法框架，旨在提供各种强大的预训练主干网络， 并支持了不同的预训练策略。MMPretrain 源自著名的开源项目 MMClassification 和 MMSelfSup，并开发了许多令人兴奋的新功能。 目前，预训练阶段对于视觉识别至关重要，凭借丰富而强大的预训练模型，我们能够改进各种下游视觉任务。

我们的代码库旨在成为一个易于使用和用户友好的代码库库，并简化学术研究活动和工程任务。 我们在以下不同部分中详细介绍了 MMPretrain 的特性和设计。

代码仓库：https://github.com/open-mmlab/mmpretrain
文档教程：https://mmpretrain.readthedocs.io/en/latest/

学习配置文件

MMPretrain 主要使用 python 文件作为配置文件，所有配置文件都放置在 configs文件夹下，目录结构如下所示：

MMPretrain/
    ├── configs/
    │   ├── _base_/                       # primitive configuration folder
    │   │   ├── datasets/                      # primitive datasets
    │   │   ├── models/                        # primitive models
    │   │   ├── schedules/                     # primitive schedules
    │   │   └── default_runtime.py             # primitive runtime setting
    │   ├── beit/                         # BEiT Algorithms Folder
    │   ├── mae/                          # MAE Algorithms Folder
    │   ├── mocov2/                       # MoCoV2 Algorithms Folder
    │   ├── resnet/                       # ResNet Algorithms Folder
    │   ├── swin_transformer/             # Swin Algorithms Folder
    │   ├── vision_transformer/           # ViT Algorithms Folder
    │   ├── ...
    └── ...

可以使用 python tools/misc/print_config.py /PATH/TO/CONFIG 命令来查看完整的配置信息，从而方便检查所对应的配置文件。

配置文件结构

在 configs/base 文件夹下有 4 个基本组件类型，分别是：

• 模型(model)
• 数据(data)
• 训练策略(schedule)
• 运行设置(runtime)

原始配置文件

_base_ = [                                    # 此配置文件将继承所有 `_base_` 中的配置
    '../_base_/models/resnet50.py',           # 模型配置
    '../_base_/datasets/imagenet_bs32.py',    # 数据配置
    '../_base_/schedules/imagenet_bs256.py',  # 训练策略配置
    '../_base_/default_runtime.py'            # 默认运行设置
]

模型配置文件

ResNet50 的模型配置’configs/base/models/resnet50.py’

model = dict(
    type='ImageClassifier',     # 主模型类型（对于图像分类任务，使用 `ImageClassifier`）
    backbone=dict(
        type='ResNet',          # 主干网络类型
        # 除了 `type` 之外的所有字段都来自 `ResNet` 类的 __init__ 方法
        # 可查阅 https://mmpretrain.readthedocs.io/zh_CN/latest/api/generated/mmpretrain.models.backbones.ResNet.html
        depth=50,
        num_stages=4,           # 主干网络状态(stages)的数目，这些状态产生的特征图作为后续的 head 的输入。
        out_indices=(3, ),      # 输出的特征图输出索引。
        frozen_stages=-1,       # 冻结主干网的层数
        style='pytorch'),
    neck=dict(type='GlobalAveragePooling'),    # 颈网络类型
    head=dict(
        type='LinearClsHead',         # 分类颈网络类型
        # 除了 `type` 之外的所有字段都来自 `LinearClsHead` 类的 __init__ 方法
        # 可查阅 https://mmpretrain.readthedocs.io/zh_CN/latest/api/generated/mmpretrain.models.heads.LinearClsHead.html
        num_classes=1000,
        in_channels=2048,
        loss=dict(type='CrossEntropyLoss', loss_weight=1.0), # 损失函数配置信息
        topk=(1, 5),                 # 评估指标，Top-k 准确率， 这里为 top1 与 top5 准确率
    ))

数据原始配置文件

ResNet50 的数据配置 ‘configs/base/datasets/imagenet_bs32.py’

dataset_type = 'ImageNet'
# 预处理配置
data_preprocessor = dict(
    # 输入的图片数据通道以 'RGB' 顺序
    mean=[123.675, 116.28, 103.53],    # 输入图像归一化的 RGB 通道均值
    std=[58.395, 57.12, 57.375],       # 输入图像归一化的 RGB 通道标准差
    to_rgb=True,                       # 是否将通道翻转，从 BGR 转为 RGB 或者 RGB 转为 BGR
)

train_pipeline = [
    dict(type='LoadImageFromFile'),     # 读取图像
    dict(type='RandomResizedCrop', scale=224),     # 随机放缩裁剪
    dict(type='RandomFlip', prob=0.5, direction='horizontal'),   # 随机水平翻转
    dict(type='PackInputs'),         # 准备图像以及标签
]

test_pipeline = [
    dict(type='LoadImageFromFile'),     # 读取图像
    dict(type='ResizeEdge', scale=256, edge='short'),  # 缩放短边尺寸至 256px
    dict(type='CenterCrop', crop_size=224),     # 中心裁剪
    dict(type='PackInputs'),                 # 准备图像以及标签
]

# 构造训练集 dataloader
train_dataloader = dict(
    batch_size=32,                     # 每张 GPU 的 batchsize
    num_workers=5,                     # 每个 GPU 的线程数
    dataset=dict(                      # 训练数据集
        type=dataset_type,
        data_root='data/imagenet',
        ann_file='meta/train.txt',
        data_prefix='train',
        pipeline=train_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=True),   # 默认采样器
    persistent_workers=True,                             # 是否保持进程，可以缩短每个 epoch 的准备时间
)

# 构造验证集 dataloader
val_dataloader = dict(
    batch_size=32,
    num_workers=5,
    dataset=dict(
        type=dataset_type,
        data_root='data/imagenet',
        ann_file='meta/val.txt',
        data_prefix='val',
        pipeline=test_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=False),
    persistent_workers=True,
)
# 验证集评估设置，使用准确率为指标， 这里使用 topk1 以及 top5 准确率
val_evaluator = dict(type='Accuracy', topk=(1, 5))

test_dataloader = val_dataloader  # test dataloader 配置，这里直接与 val_dataloader 相同
test_evaluator = val_evaluator    # 测试集的评估配置，这里直接与 val_evaluator 相同

训练策略配置

ResNet50 的训练策略配置’configs/base/schedules/imagenet_bs256.py’

optim_wrapper = dict(
    # 使用 SGD 优化器来优化参数
    optimizer=dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001))

# 学习率参数的调整策略
# 'MultiStepLR' 表示使用多步策略来调度学习率（LR）。
param_scheduler = dict(
    type='MultiStepLR', by_epoch=True, milestones=[30, 60, 90], gamma=0.1)

# 训练的配置，迭代 100 个 epoch，每一个训练 epoch 后都做验证集评估
# 'by_epoch=True' 默认使用 `EpochBaseLoop`,  'by_epoch=False' 默认使用 `IterBaseLoop`
train_cfg = dict(by_epoch=True, max_epochs=100, val_interval=1)
# 使用默认的验证循环控制器
val_cfg = dict()
# 使用默认的测试循环控制器
test_cfg = dict()

# 通过默认策略自动缩放学习率，此策略适用于总批次大小 256
# 如果你使用不同的总批量大小，比如 512 并启用自动学习率缩放
# 我们将学习率扩大到 2 倍
auto_scale_lr = dict(base_batch_size=256)

运行设置

’configs/base/default_runtime.py’

# 默认所有注册器使用的域
default_scope = 'mmpretrain'

# 配置默认的 hook
default_hooks = dict(
    # 记录每次迭代的时间。
    timer=dict(type='IterTimerHook'),

    # 每 100 次迭代打印一次日志。
    logger=dict(type='LoggerHook', interval=100),

    # 启用默认参数调度 hook。
    param_scheduler=dict(type='ParamSchedulerHook'),

    # 每个 epoch 保存检查点。
    checkpoint=dict(type='CheckpointHook', interval=1),

    # 在分布式环境中设置采样器种子。
    sampler_seed=dict(type='DistSamplerSeedHook'),

    # 验证结果可视化，默认不启用，设置 True 时启用。
    visualization=dict(type='VisualizationHook', enable=False),
)

# 配置环境
env_cfg = dict(
   # 是否开启 cudnn benchmark
    cudnn_benchmark=False,

    # 设置多进程参数
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),

    # 设置分布式参数
    dist_cfg=dict(backend='nccl'),
)

# 设置可视化工具
vis_backends = [dict(type='LocalVisBackend')] # 使用磁盘(HDD)后端
visualizer = dict(
    type='UniversalVisualizer', vis_backends=vis_backends, name='visualizer')

# 设置日志级别
log_level = 'INFO'

# 从哪个检查点加载
load_from = None

# 是否从加载的检查点恢复训练
resume = False

使用现有模型进行推理

• list_models: 列举 MMPretrain 中所有可用模型名称
• get_model: 通过模型名称或模型配置文件获取模型
• inference_model: 使用与模型相对应任务的推理器进行推理。主要用作快速 展示。如需配置进阶用法，还需要直接使用下列推理器。

推理器:

1. ImageClassificationInferencer: 对给定图像执行图像分类。
2. ImageRetrievalInferencer: 从给定的一系列图像中，检索与给定图像最相似的图像。
3. ImageCaptionInferencer: 生成给定图像的一段描述。
4. VisualQuestionAnsweringInferencer: 根据给定的图像回答问题。
5. VisualGroundingInferencer: 根据一段描述，从给定图像中找到一个与描述对应的对象。
6. TextToImageRetrievalInferencer: 从给定的一系列图像中，检索与给定文本最相似的图像。
7. ImageToTextRetrievalInferencer: 从给定的一系列文本中，检索与给定图像最相似的文本。
8. NLVRInferencer: 对给定的一对图像和一段文本进行自然语言视觉推理（NLVR 任务）。
9. FeatureExtractor: 通过视觉主干网络从图像文件提取特征。

代码实例
列出 MMPreTrain 中的所有已支持的模型。

from mmpretrain import list_models
list_models()

list_models支持 Unix 文件名风格的模式匹配，你可以使用 ** * ** 匹配任意字符。

from mmpretrain import list_models
list_models("*convnext-b*21k")

使用推理器的 list_models方法获取对应任务可用的所有模型

from mmpretrain import ImageCaptionInferencer
ImageCaptionInferencer.list_models()

获取模型

使用 get_model 获取特定模型

from mmpretrain import get_model

model = get_model("convnext-base_in21k-pre_3rdparty_in1k")

model = get_model("convnext-base_in21k-pre_3rdparty_in1k", pretrained=True)

model = get_model("convnext-base_in21k-pre_3rdparty_in1k", pretrained="your_local_checkpoint_path")

model = get_model("convnext-base_in21k-pre_3rdparty_in1k", head=dict(num_classes=10))

model_headless = get_model("resnet18_8xb32_in1k", head=None, neck=None, backbone=dict(out_indices=(1, 2, 3)))

获得的模型是一个通常的 PyTorch Module

import torch
from mmpretrain import get_model
model = get_model('convnext-base_in21k-pre_3rdparty_in1k', pretrained=True)
x = torch.rand((1, 3, 224, 224))
y = model(x)
print(type(y), y.shape)

在给定图像上进行推理

from mmpretrain import inference_model
image = 'https://github.com/open-mmlab/mmpretrain/raw/main/demo/demo.JPEG'
# 如果你没有图形界面，请设置 `show=False`
result = inference_model('resnet50_8xb32_in1k', image, show=True)
print(result['pred_class'])

上述 inference_model接口可以快速进行模型推理，但它每次调用都需要重新初始化模型，也无法进行多个样本的推理。 因此我们需要使用推理器来进行多次调用。

from mmpretrain import ImageClassificationInferencer
image = 'https://github.com/open-mmlab/mmpretrain/raw/main/demo/demo.JPEG'
inferencer = ImageClassificationInferencer('resnet50_8xb32_in1k')
# 注意推理器的输出始终为一个结果列表，即使输入只有一个样本
result = inferencer('https://github.com/open-mmlab/mmpretrain/raw/main/demo/demo.JPEG')[0]
print(result['pred_class'])
# 你可以对多张图像进行批量推理
image_list = ['demo/demo.JPEG', 'demo/bird.JPEG'] * 16
results = inferencer(image_list, batch_size=8)
print(len(results))
print(results[1]['pred_class'])

在 CUDA 上进行推理

from mmpretrain import ImageClassificationInferencer
image = 'https://github.com/open-mmlab/mmpretrain/raw/main/demo/demo.JPEG'
config = 'configs/resnet/resnet50_8xb32_in1k.py'
checkpoint = 'https://download.openmmlab.com/mmclassification/v0/resnet/resnet50_8xb32_in1k_20210831-ea4938fc.pth'
inferencer = ImageClassificationInferencer(model=config, pretrained=checkpoint, device='cuda')
result = inferencer(image)[0]
print(result['pred_class'])

训练

python tools/train.py ${CONFIG_FILE} [ARGS]

测试

python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [ARGS]