OpenMMLab Foundational Library for Training Deep Learning Models
 
 
Go to file
Mashiro 50465a2002
[Docs] Add Chinese documentation for inferencer (#884)
* [Feature] Add BaseInferencer (#773)

* Update BaseInferencer

* Fix ci

* Fix CI and rename iferencer to infer

* Fix CI

* Add renamed file

* Add test file

* Adjust interface sequence

* refine preprocess

* Update unit test

Update unit test

* Update unit test

* Fix unit test

* Fix as comment

* Minor refine

* Fix docstring and support load image from different backend

* Support load collate_fn from downstream repos, refine dispatch

* Minor refine

* Fix lint

* refine grammar

* Remove FileClient

* Refine docstring

* add rich

* Add list_models

* Add list_models

* Remove backend args

* Minor refine

* Fix typos in docs and type hints (#787)

* [Fix] Add _inputs_to_list (#795)

* Add preprocess inputs

* Add type hint

* update api/infer in index.rst

* rename preprocess_inputs to _inputs_to_list

* Fix doc format

* Update infer.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* [Fix] Fix alias type (#801)

* [Enhance] Support loading model config from checkpoint (#864)

* first commit

* [Enhance] Support build model from weight

* minor refine

* Fix type hint

* refine comments

* Update docstring

* refine as comment

* Add  method

* Refine docstring

* Fix as comment

* refine comments

* Refine warning message

* Fix unit test and refine comments

* add infer.md

* minor refine

* minor refine

* minor refine

* Apply suggestions from code review

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Fix conflicts and minor refine

* minor refine

* Fix as comment

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>
Co-authored-by: Tong Gao <gaotongxiao@gmail.com>
2023-01-16 16:31:11 +08:00
.circleci Deprecate the support for python3.6 (#863) 2023-01-06 12:50:36 +08:00
.github Deprecate the support for python3.6 (#863) 2023-01-06 12:50:36 +08:00
docker [Feature]: Add dockerfile (#347) 2022-09-07 20:26:17 +08:00
docs [Docs] Add Chinese documentation for inferencer (#884) 2023-01-16 16:31:11 +08:00
examples Make TTAModel compatible with FSDP (#611) 2022-12-27 21:39:43 +08:00
mmengine [Docs] Add Chinese documentation for inferencer (#884) 2023-01-16 16:31:11 +08:00
requirements [Feature] Add BaseInferencer (#874) 2023-01-16 16:01:06 +08:00
tests [Feature] Add BaseInferencer (#874) 2023-01-16 16:01:06 +08:00
.gitignore [Docs] Make API of docs clearer (#469) 2022-08-30 16:27:20 +08:00
.owners.yml Add issue or pr assignees (#608) 2022-10-15 14:37:04 +08:00
.pre-commit-config-zh-cn.yaml Update the repo url of pyupgrade pre-commit hook (#756) 2022-11-23 20:11:53 +08:00
.pre-commit-config.yaml [Docs] Fix Chinese docs whitespaces (#521) 2022-10-17 14:20:00 +08:00
.readthedocs.yml [Docs] Add readthedocs.yml (#501) 2022-08-31 23:34:25 +08:00
CODEOWNERS Add CODEOWNERS file to assign reviewers automactically (#605) 2022-10-13 18:51:28 +08:00
CONTRIBUTING.md [Docs] Add contributing.md and code_style.md (#754) 2022-12-04 10:52:11 +08:00
CONTRIBUTING_zh-CN.md [Docs] Add contributing.md and code_style.md (#754) 2022-12-04 10:52:11 +08:00
LICENSE
MANIFEST.in [Fix] Add torchvision_0.12.json (#820) 2022-12-13 19:13:14 +08:00
README.md [Fix] Unloaded weights will not be initialized when using PretrainedIinit (#764) 2023-01-09 18:46:30 +08:00
README_zh-CN.md [Fix] Unloaded weights will not be initialized when using PretrainedIinit (#764) 2023-01-09 18:46:30 +08:00
pytest.ini Update pytest.ini (#156) 2022-03-30 23:26:12 +08:00
requirements.txt Fix setup.py and add requirements.txt (#37) 2022-02-19 13:43:36 +08:00
setup.cfg [Dev]: Upgrade pre commit hooks (#576) 2022-10-08 10:31:19 +08:00
setup.py Deprecate the support for python3.6 (#863) 2023-01-06 12:50:36 +08:00

README.md

 
OpenMMLab website HOT      OpenMMLab platform TRY IT OUT
 

PyPI - Python Version PyPI license open issues issue resolution

🤔Reporting Issues

English | 简体中文

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.

What's New

v0.4.0 was released in 2022-12-28.

Highlights:

  • Registry supports importing modules automatically
  • Upgrade the documentation and provide the English documentation
  • Provide ProfileHook to profile the training/testing process

Read Changelog for more details.

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.

Build Models

First, we need to define a model which 1) inherits from BaseModel and 2) accepts an additional argument mode in the forward method, in addition to those arguments related to the dataset.

  • During training, the value of mode is "loss," and the forward method should return a dict containing the key "loss".
  • During validation, the value of mode is "predict", and the forward method should return results containing both predictions and labels.
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
Build Datasets

Next, we need to create Datasets and DataLoaders for training and validation. In this case, we simply use built-in datasets supported in 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)
                                ])))
Build Metrics

To validate and test the model, we need to define a Metric called accuracy to evaluate the model. This metric needs to inherit from BaseMetric and implements the process and compute_metrics methods.

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)
Build a Runner

Finally, we can construct a Runner with previously defined Model, DataLoader, and Metrics, with some other configs, as shown below.

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

runner = Runner(
    model=MMResNet50(),
    work_dir='./work_dir',
    train_dataloader=train_dataloader,
    # a wapper 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),
)
Launch Training
runner.train()

Learn More

Tutorials
Advanced tutorials
Examples
Design
Migration guide

Contributing

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

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.
  • MMClassification: OpenMMLab image classification toolbox and benchmark.
  • MMDetection: OpenMMLab detection toolbox and benchmark.
  • MMDetection3D: OpenMMLab's next-generation platform for general 3D object detection.
  • MMRotate: OpenMMLab rotated object detection toolbox and benchmark.
  • MMYOLO: OpenMMLab YOLO series toolbox and benchmark.
  • MMSegmentation: OpenMMLab semantic segmentation toolbox and benchmark.
  • MMOCR: OpenMMLab text detection, recognition, and understanding toolbox.
  • MMPose: OpenMMLab pose estimation toolbox and benchmark.
  • MMHuman3D: OpenMMLab 3D human parametric model toolbox and benchmark.
  • MMSelfSup: OpenMMLab self-supervised learning toolbox and benchmark.
  • MMRazor: OpenMMLab model compression toolbox and benchmark.
  • MMFewShot: OpenMMLab fewshot learning toolbox and benchmark.
  • MMAction2: OpenMMLab's next-generation action understanding toolbox and benchmark.
  • MMTracking: OpenMMLab video perception toolbox and benchmark.
  • MMFlow: OpenMMLab optical flow toolbox and benchmark.
  • MMEditing: OpenMMLab image and video editing toolbox.
  • MMGeneration: OpenMMLab image and video generative models toolbox.
  • MMDeploy: OpenMMLab model deployment framework.