mmengine
v0.10.5

簡介|安裝|入門|文檔|報告問題
英語| 簡體中文











V0.10.5於2024-9-11發行。
亮點:
artifact_location #1505exclude_frozen_parameters DeepSpeedEngine._zero3_consolidated_16bit_state_dict #1517閱讀ChangElog以獲取更多詳細信息。
Mmengine是一個基礎圖書館,用於培訓基於Pytorch的深度學習模型。它是所有OpenMMLAB代碼庫的培訓引擎,該引擎支持各個研究領域的數百種算法。此外,MMENGINE也適用於非openmmlab項目。它的亮點如下:
整合主流大型模型培訓框架
支持各種培訓策略
提供用戶友好的配置系統
涵蓋主流培訓監控平台
| mmengine | Pytorch | Python |
|---|---|---|
| 主要的 | > = 1.6 <= 2.1 | > = 3.8,<= 3.11 |
| > = 0.9.0,<= 0.10.4 | > = 1.6 <= 2.1 | > = 3.8,<= 3.11 |
在安裝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行的代碼中構建完整,可配置的培訓和驗證過程。
首先,我們需要定義一個模型,該BaseModel 1)除了與數據集相關的那些參數外,還接受了forward方法中的其他參數mode 。
mode的價值是“損失”, forward方法應返回包含密鑰“損失”的dict 。mode的值為“預測”,正向方法應返回包含預測和標籤的結果。 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接下來,我們需要創建數據集S和數據加載器進行培訓和驗證。在這種情況下,我們只需使用火炬手支持的內置數據集即可。
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
# 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 )最後,我們可以使用先前定義的Model , DataLoader和Metrics構建跑步者,以及其他一些配置,如下所示。
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 ()我們感謝為改善Mmengine的所有貢獻。請參閱貢獻指南的貢獻。
如果您發現此項目對您的研究有用,請考慮引用:
@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}
}
該項目以Apache 2.0許可發布。