mmengine

首先，我们需要定义一个模型，该模型 1）继承自 BaseModel，2）在 forward 方法中接受一个额外的 mode 参数，除了与数据集相关的参数之外。

• 在训练过程中，mode 的值为 "loss"，forward 方法应返回一个包含 "loss" 键的 dict。
• 在验证过程中，mode 的值为 "predict"，forward 方法应返回包含预测结果和标签的结果。

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

接下来，我们需要为训练和验证创建数据集和数据加载器。 在这个例子中，我们简单地使用 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)
                                ])))

为了验证和测试模型，我们需要定义一个名为 accuracy 的指标来评估模型。这个指标需要继承自 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)

最后，我们可以用之前定义的 Model、DataLoader 和 Metrics，以及一些其他配置来构建一个 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)),
    # 设置一些训练配置，如 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()

• Runner
• 数据集和数据加载器
• 模型
• 评估
• 优化器包装器
• 参数调度器
• 钩子

• 训练 GAN

• 恢复训练
• 加速训练
• 节省 GPU 内存

• 钩子
• 运行器
• 评估
• 可视化
• 日志记录
• 推理

• 将 Runner 从 MMCV 迁移到 MMEngine
• 将 Hook 从 MMCV 迁移到 MMEngine
• 将模型从 MMCV 迁移到 MMEngine
• 将参数调度器从 MMCV 迁移到 MMEngine
• 将数据变换迁移到 OpenMMLab 2.0