Package: mmengine
OpenMMLab websiteHOTOpenMMLab platformTRY IT OUT📘Documentation|🛠️Installation|🤔Reporting IssuesEnglish |简体中文What's Newv0.10.3 was released on 2024-1-24.Highlights:Support installing mmengine-lite with no dependency on opencv. Refer to theInstallationfor more details.Support training withColossalAI. Refer to theTraining Large Modelsfor more detailed usages.Support gradient checkpointing. Refer to theSave Memory on GPUfor more details.Supports multiple visualization backends, includingNeptuneVisBackend,DVCLiveVisBackendandAimVisBackend. Refer toVisualization Backendsfor more details.ReadChangelogfor more details.Table of ContentsIntroductionInstallationGet StartedLearn MoreContributingCitationLicenseEcosystemProjects in OpenMMLabIntroductionMMEngine 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: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.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.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.InstallationBefore installing MMEngine, please ensure that PyTorch has been successfully installed following theofficial guide.Install MMEnginepipinstall-Uopenmim
miminstallmmengineVerify the installationpython-c'from mmengine.utils.dl_utils import collect_env;print(collect_env())'Get StartedTaking 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 ModelsFirst, we need to define amodelwhich 1) inherits fromBaseModeland 2) accepts an additional argumentmodein theforwardmethod, in addition to those arguments related to the dataset.During training, the value ofmodeis "loss", and theforwardmethod should return adictcontaining the key "loss".During validation, the value ofmodeis "predict", and the forward method should return results containing both predictions and labels.importtorch.nn.functionalasFimporttorchvisionfrommmengine.modelimportBaseModelclassMMResNet50(BaseModel):def__init__(self):super().__init__()self.resnet=torchvision.models.resnet50()defforward(self,imgs,labels,mode):x=self.resnet(imgs)ifmode=='loss':return{'loss':F.cross_entropy(x,labels)}elifmode=='predict':returnx,labelsBuild DatasetsNext, we need to createDatasets andDataLoaders for training and validation.
In this case, we simply use built-in datasets supported in TorchVision.importtorchvision.transformsastransformsfromtorch.utils.dataimportDataLoadernorm_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 MetricsTo validate and test the model, we need to define aMetriccalled accuracy to evaluate the model. This metric needs to inherit fromBaseMetricand implements theprocessandcompute_metricsmethods.frommmengine.evaluatorimportBaseMetricclassAccuracy(BaseMetric):defprocess(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(),})defcompute_metrics(self,results):total_correct=sum(item['correct']foriteminresults)total_size=sum(item['batch_size']foriteminresults)# Returns a dictionary with the results of the evaluated metrics,# where the key is the name of the metricreturndict(accuracy=100*total_correct/total_size)Build a RunnerFinally, we can construct aRunnerwith previously definedModel,DataLoader, andMetrics, with some other configs, as shown below.fromtorch.optimimportSGDfrommmengine.runnerimportRunnerrunner=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 epochstrain_cfg=dict(by_epoch=True,max_epochs=5,val_interval=1),val_dataloader=val_dataloader,val_cfg=dict(),val_evaluator=dict(type=Accuracy),)Launch Trainingrunner.train()Learn MoreTutorialsRunnerDataset and DataLoaderModelEvaluationOptimWrapperParameter SchedulerHookAdvanced tutorialsRegistryConfigBaseDatasetData TransformWeight InitializationVisualizationAbstract Data ElementDistribution CommunicationLoggingFile IOGlobal manager (ManagerMixin)Use modules from other librariesTest Time AgumentationExamplesTrain a GANCommon UsageResume TrainingSpeed up TrainingSave Memory on GPUDesignHookRunnerEvaluationVisualizationLoggingInferMigration guideMigrate Runner from MMCV to MMEngineMigrate Hook from MMCV to MMEngineMigrate Model from MMCV to MMEngineMigrate Parameter Scheduler from MMCV to MMEngineMigrate Data Transform to OpenMMLab 2.0ContributingWe appreciate all contributions to improve MMEngine. Please refer toCONTRIBUTING.mdfor the contributing guideline.CitationIf 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}
}LicenseThis project is released under theApache 2.0 license.EcosystemAPES: Attention-based Point Cloud Edge SamplingDiffEngine: diffusers training toolbox with mmengineProjects in OpenMMLabMIM: 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: OpenMMLabAdvanced,Generative andIntelligentCreation 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.