MMEngine 执行器

MMEngine 执行器#

from pathlib import Path
import sys
temp_dir = Path(".temp")
sys.path.extend([str(temp_dir), str(temp_dir.resolve().parents[3]/"src")])
temp_dir.mkdir(exist_ok=True)
(temp_dir/"models").mkdir(exist_ok=True, parents=True)
import torch.nn.functional as F
import torch
from torch import nn
import torchvision
from mmengine.model import BaseModel
from mmengine.evaluator import BaseMetric
from mmengine.registry import MODELS, DATASETS, METRICS

@MODELS.register_module()
class MMResNet50(BaseModel):
    def __init__(self, data_preprocessor: dict|nn.Module|None = None,
                 init_cfg: dict|None = None):
        super().__init__(data_preprocessor=data_preprocessor, init_cfg=init_cfg)
        self.resnet = torchvision.models.resnet50()

    def forward(self, inputs: torch.Tensor,
                data_samples: list|list = None,
                mode: str = 'tensor') -> dict[str, torch.Tensor] | list:
        x = self.resnet(inputs)
        if mode == 'loss':
            return {'loss': F.cross_entropy(x, data_samples)}
        elif mode == 'predict':
            return x, data_samples
        else:
            return x
/media/pc/data/lxw/envs/anaconda3x/envs/xxx/lib/python3.12/site-packages/mmengine/optim/optimizer/zero_optimizer.py:11: DeprecationWarning: `TorchScript` support for functional optimizers is deprecated and will be removed in a future PyTorch release. Consider using the `torch.compile` optimizer instead.
  from torch.distributed.optim import \
from testing import ACTIVATION
print(ACTIVATION.module_dict)
ACTIVATION.build(dict(type='f', x=2))