mmedit.registry 源代码
# Copyright (c) OpenMMLab. All rights reserved.
"""Registries and utilities in MMEditing.
MMEditing provides 17 registry nodes to support using modules across projects.
Each node is a child of the root registry in MMEngine.
More details can be found at
https://mmengine.readthedocs.io/en/latest/tutorials/registry.html.
"""
from mmengine.registry import DATA_SAMPLERS as MMENGINE_DATA_SAMPLERS
from mmengine.registry import DATASETS as MMENGINE_DATASETS
from mmengine.registry import EVALUATOR as MMENGINE_EVALUATOR
from mmengine.registry import HOOKS as MMENGINE_HOOKS
from mmengine.registry import LOG_PROCESSORS as MMENGINE_LOG_PROCESSORS
from mmengine.registry import LOOPS as MMENGINE_LOOPS
from mmengine.registry import METRICS as MMENGINE_METRICS
from mmengine.registry import MODEL_WRAPPERS as MMENGINE_MODEL_WRAPPERS
from mmengine.registry import MODELS as MMENGINE_MODELS
from mmengine.registry import \
OPTIM_WRAPPER_CONSTRUCTORS as MMENGINE_OPTIM_WRAPPER_CONSTRUCTORS
from mmengine.registry import OPTIM_WRAPPERS as MMENGINE_OPTIM_WRAPPERS
from mmengine.registry import OPTIMIZERS as MMENGINE_OPTIMIZERS
from mmengine.registry import PARAM_SCHEDULERS as MMENGINE_PARAM_SCHEDULERS
from mmengine.registry import \
RUNNER_CONSTRUCTORS as MMENGINE_RUNNER_CONSTRUCTORS
from mmengine.registry import RUNNERS as MMENGINE_RUNNERS
from mmengine.registry import TASK_UTILS as MMENGINE_TASK_UTILS
from mmengine.registry import TRANSFORMS as MMENGINE_TRANSFORMS
from mmengine.registry import VISBACKENDS as MMENGINE_VISBACKENDS
from mmengine.registry import VISUALIZERS as MMENGINE_VISUALIZERS
from mmengine.registry import \
WEIGHT_INITIALIZERS as MMENGINE_WEIGHT_INITIALIZERS
from mmengine.registry import Registry
__all__ = [
'RUNNERS', 'RUNNER_CONSTRUCTORS', 'LOOPS', 'HOOKS', 'LOG_PROCESSORS',
'OPTIMIZERS', 'OPTIM_WRAPPERS', 'OPTIM_WRAPPER_CONSTRUCTORS',
'PARAM_SCHEDULERS', 'DATASETS', 'DATA_SAMPLERS', 'TRANSFORMS', 'MODELS',
'MODEL_WRAPPERS', 'WEIGHT_INITIALIZERS', 'TASK_UTILS',
'DIFFUSION_SCHEDULERS', 'METRICS', 'EVALUATORS', 'VISUALIZERS',
'VISBACKENDS'
]
#######################################################################
# mmedit.engine #
#######################################################################
# Runners like `EpochBasedRunner` and `IterBasedRunner`
)
# Runner constructors that define how to initialize runners
[文档]RUNNER_CONSTRUCTORS = Registry(
'runner constructor',
parent=MMENGINE_RUNNER_CONSTRUCTORS,
locations=['mmedit.engine'],
)
# Loops which define the training or test process, like `EpochBasedTrainLoop`
)
# Hooks to add additional functions during running, like `CheckpointHook`
)
# Log processors to process the scalar log data.
[文档]LOG_PROCESSORS = Registry(
'log processor',
parent=MMENGINE_LOG_PROCESSORS,
locations=['mmedit.engine'],
)
# Optimizers to optimize the model weights, like `SGD` and `Adam`.
)
# Optimizer wrappers to enhance the optimization process.
[文档]OPTIM_WRAPPERS = Registry(
'optimizer_wrapper',
parent=MMENGINE_OPTIM_WRAPPERS,
locations=['mmedit.engine'],
)
# Optimizer constructors to customize the hyper-parameters of optimizers.
[文档]OPTIM_WRAPPER_CONSTRUCTORS = Registry(
'optimizer wrapper constructor',
parent=MMENGINE_OPTIM_WRAPPER_CONSTRUCTORS,
locations=['mmedit.engine'],
)
# Parameter schedulers to dynamically adjust optimization parameters.
[文档]PARAM_SCHEDULERS = Registry(
'parameter scheduler',
parent=MMENGINE_PARAM_SCHEDULERS,
locations=['mmedit.engine'],
)
#######################################################################
# mmedit.datasets #
#######################################################################
# Datasets like `ImageNet` and `CIFAR10`.
)
# Samplers to sample the dataset.
[文档]DATA_SAMPLERS = Registry(
'data sampler',
parent=MMENGINE_DATA_SAMPLERS,
locations=['mmedit.datasets'],
)
# Transforms to process the samples from the dataset.
[文档]TRANSFORMS = Registry(
'transform',
parent=MMENGINE_TRANSFORMS,
locations=['mmedit.datasets.transforms'],
)
#######################################################################
# mmedit.models #
#######################################################################
# Neural network modules inheriting `nn.Module`.
)
# Model wrappers like 'MMDistributedDataParallel'
[文档]MODEL_WRAPPERS = Registry(
'model_wrapper',
parent=MMENGINE_MODEL_WRAPPERS,
locations=['mmedit.models'],
)
# Weight initialization methods like uniform, xavier.
[文档]WEIGHT_INITIALIZERS = Registry(
'weight initializer',
parent=MMENGINE_WEIGHT_INITIALIZERS,
locations=['mmedit.models'],
)
# Task-specific modules like anchor generators and box coders
)
# modules for diffusion models that support adding noise and denoising
[文档]DIFFUSION_SCHEDULERS = Registry(
'diffusion scheduler',
locations=['mmedit.models.diffusion_schedulers'],
)
#######################################################################
# mmedit.evaluation #
#######################################################################
# Metrics to evaluate the model prediction results.
)
# Evaluators to define the evaluation process.
)
#######################################################################
# mmedit.visualization #
#######################################################################
# Visualizers to display task-specific results.
[文档]VISUALIZERS = Registry(
'visualizer',
parent=MMENGINE_VISUALIZERS,
locations=['mmedit.visualization'],
)
# Backends to save the visualization results, like TensorBoard, WandB.
[文档]VISBACKENDS = Registry(
'vis_backend',
parent=MMENGINE_VISBACKENDS,
locations=['mmedit.visualization'],
)