open-mmlab/mmengine

 

OpenMMLab website ^HOT      OpenMMLab platform ^{TRY IT OUT}

 

📘Documentation | 🛠️Installation | 🤔Reporting Issues

What's New

v0.8.0 was released on 2023-06-30.

Highlights:

• Support training with FSDP and DeepSpeed. Refer to the example for more detailed usages.

• Introduce the pure Python style configuration file:

  • Support navigating to base configuration file in IDE
  • Support navigating to base variable in IDE
  • Support navigating to source code of class in IDE
  • Support inheriting two configuration files containing the same field
  • Load the configuration file without other third-party requirements

  Refer to the tutorial for more detailed usages.

  

Read Changelog for more details.

Introduction

MMEngine is a foundational library for training deep learning models based on PyTorch. It provides a solid engineering foundation and frees developers from writing redundant codes on workflows. It serves as the training engine of all OpenMMLab codebases, which support hundreds of algorithms in various research areas. Moreover, MMEngine is also generic to be applied to non-OpenMMLab projects.

Major features:

1. A universal and powerful runner:

   • Supports training different tasks with a small amount of code, e.g., ImageNet can be trained with only 80 lines of code (400 lines of the original PyTorch example).
   • Easily compatible with models from popular algorithm libraries such as TIMM, TorchVision, and Detectron2.

2. Open architecture with unified interfaces:

   • Handles different algorithm tasks with unified APIs, e.g., implement a method and apply it to all compatible models.
   • Provides a unified abstraction for upper-level algorithm libraries, which supports various back-end devices such as Nvidia CUDA, Mac MPS, AMD, MLU, and more for model training.

3. Customizable training process:

   • Defines the training process just like playing with Legos.
   • Provides rich components and strategies.
   • Complete controls on the training process with different levels of APIs.

Installation

Before installing MMEngine, please ensure that PyTorch has been successfully installed following the official guide.

Install MMEngine

pip install -U openmim
mim install mmengine

Verify the installation

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

Get Started

Taking the training of a ResNet-50 model on the CIFAR-10 dataset as an example, we will use MMEngine to build a complete, configurable training and validation process in less than 80 lines of code.

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
        # Save the results 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)
        # Returns a dictionary with the results of the evaluated metrics,
        # where the key is the name of the metric
        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,
    # a wrapper to execute back propagation and gradient update, etc.
    optim_wrapper=dict(optimizer=dict(type=SGD, lr=0.001, momentum=0.9)),
    # set some training configs like epochs
    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()

Learn More

Contributing

We appreciate all contributions to improve MMEngine. Please refer to CONTRIBUTING.md for the contributing guideline.

Citation

If you find this project useful in your research, please consider cite:

@article{mmengine2022,
  title   = {{MMEngine}: OpenMMLab Foundational Library for Training Deep Learning Models},
  author  = {MMEngine Contributors},
  howpublished = {\url{https://github.com/open-mmlab/mmengine}},
  year={2022}
}

License

This project is released under the Apache 2.0 license.

Projects in OpenMMLab

• MIM: MIM installs OpenMMLab packages.
• MMCV: OpenMMLab foundational library for computer vision.
• MMEval: A unified evaluation library for multiple machine learning libraries.
• MMPreTrain: OpenMMLab pre-training toolbox and benchmark.
• MMagic: OpenMMLab Advanced, Generative and Intelligent Creation toolbox.
• MMDetection: OpenMMLab detection toolbox and benchmark.
• MMYOLO: OpenMMLab YOLO series toolbox and benchmark.
• MMDetection3D: OpenMMLab's next-generation platform for general 3D object detection.
• MMRotate: OpenMMLab rotated object detection toolbox and benchmark.
• MMTracking: OpenMMLab video perception toolbox and benchmark.
• MMPose: OpenMMLab pose estimation toolbox and benchmark.
• MMSegmentation: OpenMMLab semantic segmentation toolbox and benchmark.
• MMOCR: OpenMMLab text detection, recognition, and understanding toolbox.
• MMHuman3D: OpenMMLab 3D human parametric model toolbox and benchmark.
• MMSelfSup: OpenMMLab self-supervised learning toolbox and benchmark.
• MMFewShot: OpenMMLab fewshot learning toolbox and benchmark.
• MMAction2: OpenMMLab's next-generation action understanding toolbox and benchmark.
• MMFlow: OpenMMLab optical flow toolbox and benchmark.
• MMDeploy: OpenMMLab model deployment framework.
• MMRazor: OpenMMLab model compression toolbox and benchmark.
• Playground: A central hub for gathering and showcasing amazing projects built upon OpenMMLab.