MMEngine

介绍

MMEngine 是一个基于 PyTorch 训练深度学习模型的基础库。它支持在 Linux、Windows 和 macOS 上运行。它具有以下三个特点：

通用且强大的执行器：

支持用最少的代码训练不同的任务，例如仅用 80 行代码训练 ImageNet（原始 PyTorch 示例需要 400 行）。

轻松兼容 TIMM、TorchVision 和 Detectron2 等流行算法库中的模型。

具有统一接口的开放架构：

使用统一的 API 处理不同的任务：您可以实现一次方法并将其应用于所有兼容的模型。

通过简单的高级抽象支持各种后端设备。目前，MMEngine支持在Nvidia CUDA、Mac MPS、AMD、MLU等设备上进行模型训练。

可定制的培训流程：

定义了具有“乐高”式可组合性的高度模块化训练引擎。

提供丰富的组件和策略集。

使用不同级别的 API 完全控制培训过程。

结构

上图说明了 OpenMMLab 2.0 中 MMEngine 的层次结构。MMEngine为OpenMMLab算法库实现了下一代训练架构，为OpenMMLab内30多个算法库提供了统一的执行基础。其核心组件包括训练引擎、评估引擎和模块管理。

核心模块及相关组件

训练引擎的核心模块是 Runner. Runner负责执行训练、测试和推理任务，并管理这些过程中所需的各种组件。在训练、测试和推理任务执行过程中的特定位置，设置RunnerHooks 以允许用户扩展、插入和执行自定义逻辑。主要Runner 调用以下组件来完成训练和推理循环：

数据集：负责在训练、测试和推理任务中构建数据集，并将数据输入模型。在使用中，它由 PyTorch DataLoader 包装，该加载器启动多个子进程来加载数据。

Model：接受数据并输出训练过程中的损失；在测试和推理任务期间接受数据并执行预测。在分布式环境中，模型由模型包装器（例如，MMDistributedDataParallel）包装。

Optimizer Wrapper优化器包装器：优化器包装器在训练过程中执行反向传播来优化模型，并通过统一的接口支持混合精度训练和梯度累积。

Parameter Scheduler参数调度器：在训练过程中动态调整优化器超参数，例如学习率和动量。

在训练间隔或测试阶段，指标和评估器负责评估模型的性能。Evaluator根据数据集评估模型的预测。在Evaluator 中，有一个名为 Metrics的抽象 ，它计算各种指标，例如召回率、准确性等。

为了保证接口的统一，OpenMMLab 2.0内的评估器、模型和各个算法库中的数据之间的通信接口都使用 Data Elements进行封装。

在训练和推理执行期间，上述组件可以利用日志管理模块和可视化器进行结构化和非结构化日志存储和可视化。Logging Modules：负责管理Runner执行过程中产生的各种日志信息。消息中心实现组件、运行器和日志处理器之间的数据共享，而日志处理器则处理日志信息。处理后的日志发送至Logger和Visualizer进行管理和显示。负责 Visualizer可视化模型的特征图、预测结果以及训练过程中生成的结构化日志。它支持多种可视化后端，例如 TensorBoard 和 WanDB。

通用基础模块

MMEngine还实现了算法模型执行过程中所需的各种通用基础模块，包括：
Config：在OpenMMLab算法库中，用户可以通过编写配置文件（config）来配置训练、测试过程以及相关组件。

注册中心：负责管理算法库内具有类似功能的模块。基于算法库模块的抽象，MMEngine定义了一组根注册表。算法库内的注册表可以继承自这些根注册表，从而实现跨算法库模块调用和交互。这允许在 OpenMMLab 框架内跨不同算法无缝集成和利用模块。

文件I/O：为各个模块中的文件读写操作提供统一的接口，一致支持多种文件后端系统和格式，具有可扩展性。

分布式通信原语：在分布式程序执行期间处理不同进程之间的通信。该接口抽象了分布式和非分布式环境之间的差异，并自动处理数据设备和通信后端。

其他实用程序：还有实用程序模块，例如ManagerMixin，它实现了创建和访问全局变量的方法。许多全局可访问对象的基类Runner是ManagerMixin.

15 分钟即可开始使用 MMENGINE

在本教程中，我们将以在 CIFAR-10 数据集上训练 ResNet-50 模型为例。我们将仅用 80 行代码构建一个完整且可配置的训练和验证管道MMEgnine。整个过程包括以下步骤：
建立模型

构建数据集和 DataLoader

建立评估指标

构建 Runner 并运行任务

1、建立模型

首先，我们需要建立一个模型。在MMEngine中，模型应该继承自BaseModel. 除了表示数据集输入的参数之外，其forward方法还需要接受一个名为mode 的额外参数：

对于训练， mode的值为“loss”，并且该forward方法应返回包含键“loss”的dict。
对于验证，mode 的值为“predict”，forward 方法应返回包含预测和标签的结果。

import torch.nn.functional as F
import torchvision
from mmengine.model import BaseModel


class MMResNet50(BaseModel):
    def __init__(self):
        super().__init__()
        self.resnet = torchvision.models.resnet50()

    def forward(self, imgs, labels, mode):
        x = self.resnet(imgs)
        if mode == 'loss':
            return {'loss': F.cross_entropy(x, labels)}
        elif mode == 'predict':
            return x, labels

2、构建数据集和 DataLoader

接下来，我们需要创建用于训练和验证的数据集和数据加载器。对于基本训练和验证，我们可以简单地使用 TorchVision 支持的内置数据集。

import torchvision.transforms as transforms
from torch.utils.data import DataLoader

norm_cfg = dict(mean=[0.491, 0.482, 0.447], std=[0.202, 0.199, 0.201])
train_dataloader = DataLoader(batch_size=32,
                              shuffle=True,
                              dataset=torchvision.datasets.CIFAR10(
                                  'data/cifar10',
                                  train=True,
                                  download=True,
                                  transform=transforms.Compose([
                                      transforms.RandomCrop(32, padding=4),
                                      transforms.RandomHorizontalFlip(),
                                      transforms.ToTensor(),
                                      transforms.Normalize(**norm_cfg)
                                  ])))

val_dataloader = DataLoader(batch_size=32,
                            shuffle=False,
                            dataset=torchvision.datasets.CIFAR10(
                                'data/cifar10',
                                train=False,
                                download=True,
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize(**norm_cfg)
                                ])))

3、建立评估指标

为了验证和测试模型，我们需要定义一个称为准确性的指标来评估模型。该指标需要继承BaseMetric并实现process和compute_metrics方法，其中该process方法接受数据集的输出和其他输出当mode=“predict”。该场景的输出数据是一批数据。处理完这批数据后，我们将信息保存到self.results属性中。 compute_metrics接受一个results参数。results的输入compute_metrics是保存的所有信息process（在分布式环境中，是从所有进程中results收集的信息）。使用这些信息来计算并返回dict保存评估指标的结果

To validate and test the model, we need to define a Metric called accuracy to evaluate the model. This metric needs inherit from BaseMetric and implements the process and compute_metrics methods where the process method accepts the output of the dataset and other outputs when mode=“predict”. The output data at this scenario is a batch of data. After processing this batch of data, we save the information to self.results property. compute_metrics accepts a results parameter. The input results of compute_metrics is all the information saved in process (In the case of a distributed environment, results are the information collected from all process in all the processes). Use these information to calculate and return a dict that holds the results of the evaluation metrics

from mmengine.evaluator import BaseMetric

class Accuracy(BaseMetric):
    def process(self, data_batch, data_samples):
        score, gt = data_samples
        # save the middle result of a batch to `self.results`
        self.results.append({
            'batch_size': len(gt),
            'correct': (score.argmax(dim=1) == gt).sum().cpu(),
        })

    def compute_metrics(self, results):
        total_correct = sum(item['correct'] for item in results)
        total_size = sum(item['batch_size'] for item in results)
        # return the dict containing the eval results
        # the key is the name of the metric name
        return dict(accuracy=100 * total_correct / total_size)

4、构建 Runner 并运行任务

现在我们可以使用之前定义的Model, DataLoader, and Metrics和以及其他一些配置来构建一个Runner，如下所示：ModelDataLoaderMetrics

from torch.optim import SGD
from mmengine.runner import Runner

runner = Runner(
    # the model used for training and validation.
    # Needs to meet specific interface requirements
    model=MMResNet50(),
    # working directory which saves training logs and weight files
    work_dir='./work_dir',
    # train dataloader needs to meet the PyTorch data loader protocol
    train_dataloader=train_dataloader,
    # optimize wrapper for optimization with additional features like
    # AMP, gradtient accumulation, etc
    optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)),
    # trainging coinfs for specifying training epoches, verification intervals, etc
    train_cfg=dict(by_epoch=True, max_epochs=5, val_interval=1),
    # validation dataloaer also needs to meet the PyTorch data loader protocol
    val_dataloader=val_dataloader,
    # validation configs for specifying additional parameters required for validation
    val_cfg=dict(),
    # validation evaluator. The default one is used here
    val_evaluator=dict(type=Accuracy),
)

runner.train()

最后，让我们将上面的所有代码组合成一个完整的脚本，使用MMEngine执行器进行训练和验证：

import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.optim import SGD
from torch.utils.data import DataLoader

from mmengine.evaluator import BaseMetric
from mmengine.model import BaseModel
from mmengine.runner import Runner


class MMResNet50(BaseModel):
    def __init__(self):
        super().__init__()
        self.resnet = torchvision.models.resnet50()

    def forward(self, imgs, labels, mode):
        x = self.resnet(imgs)
        if mode == 'loss':
            return {'loss': F.cross_entropy(x, labels)}
        elif mode == 'predict':
            return x, labels


class Accuracy(BaseMetric):
    def process(self, data_batch, data_samples):
        score, gt = data_samples
        self.results.append({
            'batch_size': len(gt),
            'correct': (score.argmax(dim=1) == gt).sum().cpu(),
        })

    def compute_metrics(self, results):
        total_correct = sum(item['correct'] for item in results)
        total_size = sum(item['batch_size'] for item in results)
        return dict(accuracy=100 * total_correct / total_size)


norm_cfg = dict(mean=[0.491, 0.482, 0.447], std=[0.202, 0.199, 0.201])
train_dataloader = DataLoader(batch_size=32,
                              shuffle=True,
                              dataset=torchvision.datasets.CIFAR10(
                                  'data/cifar10',
                                  train=True,
                                  download=True,
                                  transform=transforms.Compose([
                                      transforms.RandomCrop(32, padding=4),
                                      transforms.RandomHorizontalFlip(),
                                      transforms.ToTensor(),
                                      transforms.Normalize(**norm_cfg)
                                  ])))

val_dataloader = DataLoader(batch_size=32,
                            shuffle=False,
                            dataset=torchvision.datasets.CIFAR10(
                                'data/cifar10',
                                train=False,
                                download=True,
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize(**norm_cfg)
                                ])))

runner = Runner(
    model=MMResNet50(),
    work_dir='./work_dir',
    train_dataloader=train_dataloader,
    optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)),
    train_cfg=dict(by_epoch=True, max_epochs=5, val_interval=1),
    val_dataloader=val_dataloader,
    val_cfg=dict(),
    val_evaluator=dict(type=Accuracy),
)
runner.train()

常见用法

恢复培训

加快训练速度

混合精准训练

Nvidia 将 Tensor Core 单元引入 Volta 和 Turing 架构中，以支持 FP32 和 FP16 混合精度计算。它们进一步支持 Ampere 架构中的 BF16。启用自动混合精度训练后，部分算子运行在 FP16/BF16 下，其余算子运行在 FP32 上，在不改变模型或降低训练精度的情况下，减少了训练时间和存储需求，从而支持更大批量、更大模型的训练，更大的输入尺寸。
MMEngine 提供了自动混合精度训练的包装器type=‘AmpOptimWrapper’ in optim_wrapper，只需设置即可启用自动混合精度训练，无需更改其他代码。 optim_wrapper

runner = Runner(
    model=ResNet18(),
    work_dir='./work_dir',
    train_dataloader=train_dataloader_cfg,
    optim_wrapper=dict(
        type='AmpOptimWrapper',
        # If you want to use bfloat16, uncomment the following line
        # dtype='bfloat16',  # valid values: ('float16', 'bfloat16', None)
        optimizer=dict(type='SGD', lr=0.001, momentum=0.9)),
    train_cfg=dict(by_epoch=True, max_epochs=3),
)
runner.train()

使用更快的优化器

如果使用Ascend设备，可以使用Ascend优化器来缩短模型的训练时间。Ascend设备支持的优化器如下：

NpuFusedAdadelta

NpuFusedAdam

NpuFusedAdamP

NpuFusedAdamW

NpuFusedBertAdam

NpuFusedLamb

NpuFusedRMSprop

NpuFusedRMSpropTF

NpuFusedSGD

节省 GPU 内存

内存容量对于深度学习训练和推理至关重要，决定了模型能否成功运行。常见的内存节省方法包括：

启用高效Conv BN评估特征(实验)

我们最近在 MMCV 中引入了一个实验性功能：基于本文讨论的概念的：Efficient Conv BN Eval 。设计此功能的目的是在不影响性能的情况下减少网络训练期间的内存占用。如果您的网络架构包含一系列连续的 Conv+BN 块，并且这些归一化层在训练过程中保持模式不变（使用MMDetectioneval训练目标检测器时常见），此功能可以减少内存消耗超过 $20$ 。要启用 Efficient Conv BN Eval 功能，只需添加以下命令行参数：

--cfg-options efficient_conv_bn_eval="[backbone]"

"Enabling the “efficient_conv_bn_eval” feature for these modules …在输出日志中，该功能已成功启用。由于目前正处于实验阶段，我们热切期待听到您的使用体验。请在此讨论线程中分享您的使用报告、观察结果和建议。您的反馈对于进一步开发以及确定是否应将此功能集成到稳定版本中至关重要。

梯度累积

梯度累积是一种按照配置的步数运行累积梯度而不是更新参数的机制，之后更新网络参数并清除梯度。通过这种延迟参数更新的技术，结果类似于使用大批量大小的场景，同时可以节省激活的内存。但需要注意的是，如果模型包含批量归一化层，使用梯度累积会影响性能。
配置可以这样写：

optim_wrapper_cfg = dict(
    type='OptimWrapper',
    optimizer=dict(type='SGD', lr=0.001, momentum=0.9),
    # update every four times
    accumulative_counts=4)

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from mmengine.runner import Runner
from mmengine.model import BaseModel

train_dataset = [(torch.ones(1, 1), torch.ones(1, 1))] * 50
train_dataloader = DataLoader(train_dataset, batch_size=2)


class ToyModel(BaseModel):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, img, label, mode):
        feat = self.linear(img)
        loss1 = (feat - label).pow(2)
        loss2 = (feat - label).abs()
        return dict(loss1=loss1, loss2=loss2)


runner = Runner(
    model=ToyModel(),
    work_dir='tmp_dir',
    train_dataloader=train_dataloader,
    train_cfg=dict(by_epoch=True, max_epochs=1),
    optim_wrapper=dict(optimizer=dict(type='SGD', lr=0.01),
                       accumulative_counts=4)
)
runner.train()

梯度检查点

梯度检查点是一种时间换空间方法，通过减少保存的激活数量来压缩模型，但是在计算梯度时必须重新计算未存储的激活。相应的功能已经在torch.utils.checkpoint包中实现了。其实现可以简单地总结为，在前向阶段，传递给检查点的前向函数以torch.no_grad模式运行，并且仅保存前向函数的输入和输出。然后在向后阶段重新计算其中间激活。

大模型训练技术

最近的研究表明，训练大型模型有助于提高性能，但训练如此规模的模型需要巨大的资源，并且很难将整个模型存储在单个显卡的内存中。因此引入了大型模型训练技术，例如DeepSpeed ZeRO和 FairScale 中引入的完全共享数据并行 ( FSDP ) 技术。这些技术允许在并行进程之间对参数、梯度和优化器状态进行切片，同时仍然保持数据并行性的简单性。
FSDP从 PyTorch 1.11 开始正式支持。配置可以这样写：

# located in cfg file
model_wrapper_cfg=dict(type='MMFullyShardedDataParallel', cpu_offload=True)

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from mmengine.runner import Runner
from mmengine.model import BaseModel

train_dataset = [(torch.ones(1, 1), torch.ones(1, 1))] * 50
train_dataloader = DataLoader(train_dataset, batch_size=2)


class ToyModel(BaseModel):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, img, label, mode):
        feat = self.linear(img)
        loss1 = (feat - label).pow(2)
        loss2 = (feat - label).abs()
        return dict(loss1=loss1, loss2=loss2)


runner = Runner(
    model=ToyModel(),
    work_dir='tmp_dir',
    train_dataloader=train_dataloader,
    train_cfg=dict(by_epoch=True, max_epochs=1),
    optim_wrapper=dict(optimizer=dict(type='SGD', lr=0.01)),
    cfg=dict(model_wrapper_cfg=dict(type='MMFullyShardedDataParallel', cpu_offload=True))
)
runner.train()

训练大模型

训练大型模型时，需要大量资源。单个 GPU 内存往往不足以满足训练需求。因此，开发了训练大型模型的技术，一种典型的方法是DeepSpeed ZeRO。DeepSpeed ZeRO 支持优化器、梯度和参数分片。

DeepSpeed是微软开发的一个基于PyTorch的开源分布式框架。它支持诸如ZeRO、3D-Parallelism、DeepSpeed-MoE和ZeRO-Infinity 之类的培训策略。

从 MMEngine v0.8.0 开始，MMEngine 支持使用 DeepSpeed 训练模型。

要使用 DeepSpeed，您需要先通过运行以下命令来安装它：

pip install deepspeed

安装DeepSpeed后，需要对flexiblerrunner的策略和optim_wrapper参数进行如下配置:
strategy:设置type=‘DeepSpeedStrategy’，并配置其他参数。有关更多细节，请参阅DeepSpeedStrategy。
optim_wrapper:设置type='DeepSpeedOptimwrapper’并配置其他参数。
DeepSpeedOptimWrapper for more details.
以下是与 DeepSpeed 相关的示例配置：

from mmengine.runner._flexible_runner import FlexibleRunner

# set `type='DeepSpeedStrategy'` and configure other parameters
strategy = dict(
    type='DeepSpeedStrategy',
    fp16=dict(
        enabled=True,
        fp16_master_weights_and_grads=False,
        loss_scale=0,
        loss_scale_window=500,
        hysteresis=2,
        min_loss_scale=1,
        initial_scale_power=15,
    ),
    inputs_to_half=[0],
    zero_optimization=dict(
        stage=3,
        allgather_partitions=True,
        reduce_scatter=True,
        allgather_bucket_size=50000000,
        reduce_bucket_size=50000000,
        overlap_comm=True,
        contiguous_gradients=True,
        cpu_offload=False),
)

# set `type='DeepSpeedOptimWrapper'` and configure other parameters
optim_wrapper = dict(
    type='DeepSpeedOptimWrapper',
    optimizer=dict(type='AdamW', lr=1e-3))

# construct FlexibleRunner
runner = FlexibleRunner(
    model=MMResNet50(),
    work_dir='./work_dirs',
    strategy=strategy,
    train_dataloader=train_dataloader,
    optim_wrapper=optim_wrapper,
    param_scheduler=dict(type='LinearLR'),
    train_cfg=dict(by_epoch=True, max_epochs=10, val_interval=1),
    val_dataloader=val_dataloader,
    val_cfg=dict(),
    val_evaluator=dict(type=Accuracy))

# start training
runner.train()

使用两个GPU启动分布式训练：

torchrun --nproc-per-node 2 examples/distributed_training_with_flexible_runner.py --use-deepspeed

更好的性能优化器

本文档提供了MMEngine支持的一些第三方优化器，可能会带来更快的收敛速度或者更高的性能。

D-Adaptation

D-Adaptation提供DAdaptAdaGrad,DAdaptAdam和DAdaptSGD优化器。
注意：如果您使用D-Adaptation提供的优化器，则需要将mmengine升级到0.6.0.
安装

pip install dadaptation

以DAdaptAdaGrad为例。

runner = Runner(
    model=ResNet18(),
    work_dir='./work_dir',
    train_dataloader=train_dataloader_cfg,
    # To view the input parameters for DAdaptAdaGrad, you can refer to
    # https://github.com/facebookresearch/dadaptation/blob/main/dadaptation/dadapt_adagrad.py
    optim_wrapper=dict(optimizer=dict(type='DAdaptAdaGrad', lr=0.001, momentum=0.9)),
    train_cfg=dict(by_epoch=True, max_epochs=3),
)
runner.train()

调试技巧

设置数据集的长度

在调试代码的过程中，有时需要训练几个epoch，例如调试验证过程或检查检查点保存是否符合预期。但是，如果数据集太大，可能需要很长时间才能完成一个epoch，在这种情况下可以设置数据集的长度。请注意，只有继承自BaseDataset 的数据集才支持此功能，并且 BaseDataset 的使用可以在 BaseDataset 中找到

Turn off the training and set indices as 5000 in the dataset field in configs/base/datasets/cifar10_bs16.py.

train_dataloader = dict(
    batch_size=16,
    num_workers=2,
    dataset=dict(
        type=dataset_type,
        data_prefix='data/cifar10',
        test_mode=False,
        indices=5000,  # set indices=5000，represent every epoch only iterator 5000 samples
        pipeline=train_pipeline),
    sampler=dict(type='DefaultSampler', shuffle=True),
)

再次启动训练

python tools/train.py configs/resnet/resnet18_8xb16_cifar10.py

正如我们所看到的，迭代次数已更改为313。与之前相比，这可以更快完成一个纪元。

02/20 14:44:58 - mmengine - INFO - Epoch(train)   [1][100/313]  lr: 1.0000e-01  eta: 0:31:09  time: 0.0154  data_time: 0.0004  memory: 214  loss: 2.1852
02/20 14:44:59 - mmengine - INFO - Epoch(train)   [1][200/313]  lr: 1.0000e-01  eta: 0:23:18  time: 0.0143  data_time: 0.0002  memory: 214  loss: 2.0424
02/20 14:45:01 - mmengine - INFO - Epoch(train)   [1][300/313]  lr: 1.0000e-01  eta: 0:20:39  time: 0.0143  data_time: 0.0003  memory: 214  loss: 1.814

RUNNER