OpenMMLab 官网 ^热门      OpenMMLab 平台 ^立即体验

 

简介 | 安装 | 快速入门 | 📘文档 | 🤔提交问题

最新动态

v0.10.6 于 2025年1月13日发布。

亮点：

• 支持在 MLflowVisBackend 中自定义 artifact_location #1505
• 为 DeepSpeedEngine._zero3_consolidated_16bit_state_dict 启用 exclude_frozen_parameters #1517

更多详情请参阅 Changelog。

简介

MMEngine 是一个基于 PyTorch 的深度学习模型训练基础库。它是所有 OpenMMLab 代码库的训练引擎，支持多个研究领域的数百种算法。此外，MMEngine 也可广泛应用于非 OpenMMLab 项目。其主要特点如下：

集成主流大规模模型训练框架

• ColossalAI
• DeepSpeed
• FSDP

支持多种训练策略

• 混合精度训练
• 梯度累积
• 梯度检查点

提供用户友好的配置系统

• 纯 Python 风格的配置文件，易于理解
• 纯文本风格的配置文件，支持 JSON 和 YAML

覆盖主流训练监控平台

• TensorBoard | WandB | MLflow
• ClearML | Neptune | DVCLive | Aim

安装

在安装 MMEngine 之前，请确保已按照 官方指南 成功安装 PyTorch。

安装 MMEngine

pip install -U openmim
mim install mmengine

验证安装

python -c 'from mmengine.utils.dl_utils import collect_env;print(collect_env())'

开始使用

以在 CIFAR-10 数据集上训练 ResNet-50 模型为例，我们将使用 MMEngine 在不到 80 行代码中构建一个完整且可配置的训练和验证流程。

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

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)
                                ])))

from mmengine.evaluator import BaseMetric

class Accuracy(BaseMetric):
    def process(self, data_batch, data_samples):
        score, gt = data_samples
        # 将一个批次的结果保存到 `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 dict(accuracy=100 * total_correct / total_size)

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

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()

了解更多

贡献

我们非常感谢所有对MMEngine改进的贡献。请参阅CONTRIBUTING.md以获取贡献指南。

引用

如果您在研究中发现本项目有用，请考虑引用：

@article{mmengine2022,
  title   = {{MMEngine}: OpenMMLab用于训练深度学习模型的基础库},
  author  = {MMEngine贡献者},
  howpublished = {\url{https://github.com/open-mmlab/mmengine}},
  year={2022}
}

许可证

本项目采用Apache 2.0许可证发布。

生态系统

• APES: 基于注意力的点云边缘采样
• DiffEngine: 使用mmengine的diffusers训练工具箱

OpenMMLab中的项目

• MIM: MIM用于安装OpenMMLab的相关软件包。
• MMCV: OpenMMLab计算机视觉领域的基础库。
• MMEval: 一个用于多种机器学习框架的统一评估库。
• MMPreTrain: OpenMMLab预训练工具箱及基准测试平台。
• MMagic: OpenMMLab Advanced, Generative和Intelligent Creation工具箱。
• MMDetection: OpenMMLab目标检测工具箱及基准测试平台。
• MMYOLO: OpenMMLab YOLO系列工具箱及基准测试平台。
• MMDetection3D: OpenMMLab下一代通用3D目标检测平台。
• MMRotate: OpenMMLab旋转目标检测工具箱及基准测试平台。
• MMTracking: OpenMMLab视频感知工具箱及基准测试平台。
• MMPose: OpenMMLab人体姿态估计工具箱及基准测试平台。
• MMSegmentation: OpenMMLab语义分割工具箱及基准测试平台。
• MMOCR: OpenMMLab文本检测、识别和理解工具箱。
• MMHuman3D: OpenMMLab3D人体参数化模型工具箱及基准测试平台。
• MMSelfSup: OpenMMLab自监督学习工具箱及基准测试平台。
• MMFewShot: OpenMMLab少样本学习工具箱及基准测试平台。
• MMAction2: OpenMMLab下一代动作理解工具箱及基准测试平台。
• MMFlow: OpenMMLab光流计算工具箱及基准测试平台。
• MMDeploy: OpenMMLab模型部署框架。
• MMRazor: OpenMMLab模型压缩工具箱及基准测试平台。
• Playground: 一个汇集并展示基于OpenMMLab构建的优秀项目的中心平台。

MMEngine 快速上手指南

MMEngine 是 OpenMMLab 开源的基于 PyTorch 的深度学习模型训练基础库。它提供了灵活的配置系统、主流的大模型训练框架支持（如 DeepSpeed、FSDP）以及丰富的训练策略，适用于从入门到大规模模型训练的各类场景。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, Windows, macOS
• Python: 3.8 ~ 3.11
• PyTorch: 1.6 ~ 2.1 (请先参考 PyTorch 官方指南 安装对应版本的 PyTorch)
• CUDA: 根据您的显卡驱动和 PyTorch 版本选择合适的 CUDA 版本（可选，用于 GPU 加速）

2. 安装步骤

推荐使用 openmim 工具进行安装，它能自动处理依赖关系。

步骤 1：安装 openmim

pip install -U openmim

步骤 2：安装 MMEngine

mim install mmengine

提示：如果下载速度较慢，可以添加国内镜像源加速：

pip install -U openmim -i https://pypi.tuna.tsinghua.edu.cn/simple
mim install mmengine

步骤 3：验证安装 运行以下命令检查是否安装成功及查看环境信息：

python -c 'from mmengine.utils.dl_utils import collect_env;print(collect_env())'

若无报错并输出了环境配置信息，则说明安装成功。

3. 基本使用

以下示例展示如何使用 MMEngine 在 CIFAR-10 数据集上训练一个 ResNet-50 模型。整个流程包含定义模型、数据加载、评估指标和启动训练四个步骤。

第一步：定义模型

继承 BaseModel 并重写 forward 方法，需支持 mode 参数以区分训练（计算 loss）和验证（返回预测结果）。

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

第二步：准备数据

构建训练集和验证集的 DataLoader。此处直接使用 TorchVision 内置的 CIFAR-10 数据集。

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)
        ]))
)

第三步：定义评估指标

继承 BaseMetric 实现准确率计算，需实现 process（处理单批次数据）和 compute_metrics（汇总计算）方法。

from mmengine.evaluator import BaseMetric

class Accuracy(BaseMetric):
    def process(self, data_batch, data_samples):
        score, gt = data_samples
        # 将单批次结果保存至 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 dict(accuracy=100 * total_correct / total_size)

第四步：构建 Runner 并启动训练

使用 Runner 统一管理模型、数据、优化器和训练配置，一键启动训练。

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

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()