[Feature] Support auto import modules from registry

configs/tsne/resnet50_imagenet.py

@@ -14,7 +14,7 @@
         depth=50,
         in_channels=3,
         num_stages=4,
-        out_indices=(3),
+        out_indices=(3, ),
         norm_cfg=dict(type='BN'),
         frozen_stages=-1),
     neck=dict(type='GlobalAveragePooling'),

demo/mae_visualization.ipynb

@@ -62,8 +62,7 @@
     "import torch\n",
     "from mmengine.dataset import Compose, default_collate\n",
     "\n",
-    "from mmselfsup.apis import inference_model, init_model\n",
-    "from mmselfsup.utils import register_all_modules"
+    "from mmselfsup.apis import inference_model, init_model"
    ]
   },
   {
@@ -238,7 +237,6 @@
    ],
    "source": [
     "# make random mask reproducible (comment out to make it change)\n",
-    "register_all_modules()\n",
     "torch.manual_seed(2)"
    ]
   },

demo/mmselfsup_colab_tutorial.ipynb

@@ -157,7 +157,7 @@
             ],
             "source": [
                 "# Check PyTorch installation\n",
-                "import torch, torchvision\n",
+                "import torch\n",
                 "print(torch.__version__)\n",
                 "print(torch.cuda.is_available())"
             ]
@@ -313,7 +313,7 @@
             "source": [
                 "!pip3 install openmim\n",
                 "!pip install -U openmim\n",
-                "!mim install 'mmengine' 'mmcv>=2.0.0rc1'"
+                "!mim install 'mmengine' 'mmcv>=2.0.0rc4'"
             ]
         },
         {
@@ -1418,15 +1418,8 @@
                 }
             ],
             "source": [
-                "from mmengine.config import Config, DictAction\n",
                 "from mmengine.runner import Runner\n",
                 "\n",
-                "from mmselfsup.utils import register_all_modules\n",
-                "\n",
-                "# register all modules in mmselfsup into the registries\n",
-                "# do not init the default scope here because it will be init in the runner\n",
-                "register_all_modules(init_default_scope=False)\n",
-                "\n",
                 "# build the runner from config\n",
                 "runner = Runner.from_cfg(cfg)\n",
                 "\n",
@@ -2669,15 +2662,8 @@
                 }
             ],
             "source": [
-                "from mmengine.config import Config, DictAction\n",
                 "from mmengine.runner import Runner\n",
                 "\n",
-                "from mmselfsup.utils import register_all_modules\n",
-                "\n",
-                "# register all modules in mmselfsup into the registries\n",
-                "# do not init the default scope here because it will be init in the runner\n",
-                "register_all_modules(init_default_scope=False)\n",
-                "\n",
                 "# build the runner from config\n",
                 "runner = Runner.from_cfg(benchmark_cfg)\n",
                 "\n",

mmselfsup/__init__.py

@@ -10,7 +10,7 @@
 mmengine_maximum_version = '1.0.0'
 mmengine_version = digit_version(mmengine.__version__)
 
-mmcv_minimum_version = '2.0.0rc1'
+mmcv_minimum_version = '2.0.0rc4'
 mmcv_maximum_version = '2.1.0'
 mmcv_version = digit_version(mmcv.__version__)
 

mmselfsup/apis/inference.py

@@ -6,6 +6,7 @@
 import torch
 from mmengine.config import Config
 from mmengine.dataset import Compose, default_collate
+from mmengine.registry import init_default_scope
 from mmengine.runner import load_checkpoint
 from torch import nn
 
@@ -36,8 +37,11 @@ def init_model(config: Union[str, Config],
     elif not isinstance(config, Config):
         raise TypeError('config must be a filename or Config object, '
                         f'but got {type(config)}')
+
     if options is not None:
         config.merge_from_dict(options)
+    init_default_scope(config.get('default_scope', 'mmselfsup'))
+
     config.model.pretrained = None
     config.model.setdefault('data_preprocessor',
                             config.get('data_preprocessor', None))

mmselfsup/registry.py

@@ -54,55 +54,93 @@
 
 # Registries For Runner and the related
 # manage all kinds of runners like `EpochBasedRunner` and `IterBasedRunner`
-RUNNERS = Registry('runner', parent=MMENGINE_RUNNERS)
+RUNNERS = Registry(
+    'runner', parent=MMENGINE_RUNNERS, locations=['mmselfsup.engine.runner'])
 # manage runner constructors that define how to initialize runners
 RUNNER_CONSTRUCTORS = Registry(
-    'runner constructor', parent=MMENGINE_RUNNER_CONSTRUCTORS)
+    'runner constructor',
+    parent=MMENGINE_RUNNER_CONSTRUCTORS,
+    locations=['mmselfsup.engine.runner'])
 # manage all kinds of loops like `EpochBasedTrainLoop`
-LOOPS = Registry('loop', parent=MMENGINE_LOOPS)
+LOOPS = Registry(
+    'loop', parent=MMENGINE_LOOPS, locations=['mmselfsup.engine.runner'])
 # manage all kinds of hooks like `CheckpointHook`
-HOOKS = Registry('hook', parent=MMENGINE_HOOKS)
+HOOKS = Registry(
+    'hook', parent=MMENGINE_HOOKS, locations=['mmselfsup.engine.hooks'])
 
 # Registries For Data and the related
 # manage data-related modules
-DATASETS = Registry('dataset', parent=MMENGINE_DATASETS)
-DATA_SAMPLERS = Registry('data sampler', parent=MMENGINE_DATA_SAMPLERS)
-TRANSFORMS = Registry('transform', parent=MMENGINE_TRANSFORMS)
+DATASETS = Registry(
+    'dataset', parent=MMENGINE_DATASETS, locations=['mmselfsup.datasets'])
+DATA_SAMPLERS = Registry(
+    'data sampler',
+    parent=MMENGINE_DATA_SAMPLERS,
+    locations=['mmselfsup.datasets.samplers'])
+TRANSFORMS = Registry(
+    'transform',
+    parent=MMENGINE_TRANSFORMS,
+    locations=['mmselfsup.datasets.transforms'])
 
 # manage all kinds of modules inheriting `nn.Module`
-MODELS = Registry('model', parent=MMENGINE_MODELS)
+MODELS = Registry(
+    'model', parent=MMENGINE_MODELS, locations=['mmselfsup.models'])
 # manage all kinds of model wrappers like 'MMDistributedDataParallel'
-MODEL_WRAPPERS = Registry('model_wrapper', parent=MMENGINE_MODEL_WRAPPERS)
+MODEL_WRAPPERS = Registry(
+    'model_wrapper',
+    parent=MMENGINE_MODEL_WRAPPERS,
+    locations=['mmselfsup.models'])
 # manage all kinds of weight initialization modules like `Uniform`
 WEIGHT_INITIALIZERS = Registry(
-    'weight initializer', parent=MMENGINE_WEIGHT_INITIALIZERS)
+    'weight initializer',
+    parent=MMENGINE_WEIGHT_INITIALIZERS,
+    locations=['mmselfsup.models'])
 
 # Registries For Optimizer and the related
 # manage all kinds of optimizers like `SGD` and `Adam`
-OPTIMIZERS = Registry('optimizer', parent=MMENGINE_OPTIMIZERS)
+OPTIMIZERS = Registry(
+    'optimizer',
+    parent=MMENGINE_OPTIMIZERS,
+    locations=['mmselfsup.engine.optimizers'])
 # manage optimizer wrapper
-OPTIM_WRAPPERS = Registry('optimizer_wrapper', parent=MMENGINE_OPTIM_WRAPPERS)
+OPTIM_WRAPPERS = Registry(
+    'optimizer_wrapper',
+    parent=MMENGINE_OPTIM_WRAPPERS,
+    locations=['mmselfsup.engine.optimizers'])
 # manage constructors that customize the optimization hyperparameters.
 OPTIM_WRAPPER_CONSTRUCTORS = Registry(
     'optimizer wrapper constructor',
-    parent=MMENGINE_OPTIM_WRAPPER_CONSTRUCTORS)
+    parent=MMENGINE_OPTIM_WRAPPER_CONSTRUCTORS,
+    locations=['mmselfsup.engine.optimizers'])
 # manage all kinds of parameter schedulers like `MultiStepLR`
 PARAM_SCHEDULERS = Registry(
-    'parameter scheduler', parent=MMENGINE_PARAM_SCHEDULERS)
+    'parameter scheduler',
+    parent=MMENGINE_PARAM_SCHEDULERS,
+    locations=['mmselfsup.engine.schedulers'])
 
 # manage all kinds of metrics
-METRICS = Registry('metric', parent=MMENGINE_METRICS)
+METRICS = Registry(
+    'metric', parent=MMENGINE_METRICS, locations=['mmselfsup.evaluation'])
 # manage evaluator
-EVALUATOR = Registry('evaluator', parent=MMENGINE_EVALUATOR)
+EVALUATOR = Registry(
+    'evaluator', parent=MMENGINE_EVALUATOR, locations=['mmselfsup.evaluation'])
 
 # manage task-specific modules like anchor generators and box coders
-TASK_UTILS = Registry('task util', parent=MMENGINE_TASK_UTILS)
+TASK_UTILS = Registry(
+    'task util', parent=MMENGINE_TASK_UTILS, locations=['mmselfsup.models'])
 
-# Registries For Visualizer and the related
 # manage visualizer
-VISUALIZERS = Registry('visualizer', parent=MMENGINE_VISUALIZERS)
+VISUALIZERS = Registry(
+    'visualizer',
+    parent=MMENGINE_VISUALIZERS,
+    locations=['mmselfsup.visualization'])
 # manage visualizer backend
-VISBACKENDS = Registry('vis_backend', parent=MMENGINE_VISBACKENDS)
+VISBACKENDS = Registry(
+    'vis_backend',
+    parent=MMENGINE_VISBACKENDS,
+    locations=['mmselfsup.visualization'])
 
 # manage logprocessor
-LOG_PROCESSORS = Registry('log_processor', parent=MMENGINE_LOG_PROCESSORS)
+LOG_PROCESSORS = Registry(
+    'log_processor',
+    parent=MMENGINE_LOG_PROCESSORS,
+    locations=['mmselfsup.visualization'])

tests/test_apis/test_inference.py

@@ -7,11 +7,11 @@
 import torch
 import torch.nn as nn
 from mmengine.config import Config
+from mmengine.registry import init_default_scope
 
 from mmselfsup.apis import inference_model
 from mmselfsup.models import BaseModel
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
 
 backbone = dict(
     type='ResNet',
@@ -37,12 +37,11 @@ def extract_feat(self,
 
 @pytest.mark.skipif(platform.system() == 'Windows', reason='')
 def test_inference_model():
-    register_all_modules()
-
     # Specify the data settings
     cfg = Config.fromfile(
         'configs/selfsup/relative_loc/relative-loc_resnet50_8xb64-steplr-70e_in1k.py'  # noqa: E501
     )
+    init_default_scope(cfg.get('default_scope', 'mmselfsup'))
     # Build the algorithm
     model = ExampleModel()
     model.cfg = cfg

tests/test_datasets/test_deepcluster_imagenet.py

@@ -2,9 +2,9 @@
 import os.path as osp
 
 import pytest
+from mmengine.registry import init_default_scope
 
 from mmselfsup.datasets import DeepClusterImageNet
-from mmselfsup.utils import register_all_modules
 
 # dataset settings
 train_pipeline = [
@@ -14,7 +14,7 @@
 
 
 def test_deepcluster_dataset():
-    register_all_modules()
+    init_default_scope('mmselfsup')
 
     data = dict(
         ann_file=osp.join(

tests/test_datasets/test_image_list_dataset.py

@@ -3,9 +3,9 @@
 
 import numpy as np
 import pytest
+from mmengine.registry import init_default_scope
 
 from mmselfsup.datasets import ImageList
-from mmselfsup.utils import register_all_modules
 
 # dataset settings
 train_pipeline = [
@@ -15,7 +15,7 @@
 
 
 def test_image_list_dataset():
-    register_all_modules()
+    init_default_scope('mmselfsup')
 
     data = dict(
         ann_file='',

tests/test_models/test_algorithms/test_barlowtwins.py

@@ -7,9 +7,6 @@
 
 from mmselfsup.models import BarlowTwins
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
-
-register_all_modules()
 
 backbone = dict(
     type='ResNet',

tests/test_models/test_algorithms/test_beitv1.py

@@ -7,7 +7,6 @@
 
 from mmselfsup.models import BEiT
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
 
 data_preprocessor = dict(
     type='TwoNormDataPreprocessor',
@@ -37,8 +36,6 @@
 
 @pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
 def test_beitv1():
-    register_all_modules()
-
     model = BEiT(
         backbone=backbone,
         neck=neck,

tests/test_models/test_algorithms/test_beitv2.py

@@ -7,7 +7,6 @@
 
 from mmselfsup.models import BEiT
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
 
 data_preprocessor = dict(
     type='TwoNormDataPreprocessor',
@@ -70,8 +69,6 @@
 
 @pytest.mark.skipif(platform.system() == 'Windows', reason='Windows mem limit')
 def test_beitv2():
-    register_all_modules()
-
     model = BEiT(
         backbone=backbone,
         neck=neck,

tests/test_models/test_algorithms/test_byol.py

@@ -7,9 +7,7 @@
 
 from mmselfsup.models.algorithms.byol import BYOL
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
 
-register_all_modules()
 backbone = dict(
     type='ResNet',
     depth=18,

tests/test_models/test_algorithms/test_cae.py

@@ -7,9 +7,6 @@
 
 from mmselfsup.models.algorithms import CAE
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
-
-register_all_modules()
 
 # model settings
 backbone = dict(type='CAEViT', arch='b', patch_size=16, init_values=0.1)

tests/test_models/test_algorithms/test_deepcluster.py

@@ -8,9 +8,6 @@
 
 from mmselfsup.models.algorithms import DeepCluster
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
-
-register_all_modules()
 
 num_classes = 5
 with_sobel = True,

tests/test_models/test_algorithms/test_densecl.py

@@ -9,9 +9,6 @@
 import mmselfsup
 from mmselfsup.models.algorithms.densecl import DenseCL
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
-
-register_all_modules()
 
 queue_len = 32
 feat_dim = 2

tests/test_models/test_algorithms/test_eva.py

@@ -8,9 +8,6 @@
 
 from mmselfsup.models.algorithms import EVA
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
-
-register_all_modules()
 
 backbone = dict(type='MAEViT', arch='b', patch_size=16, mask_ratio=0.75)
 neck = dict(

tests/test_models/test_algorithms/test_mae.py

@@ -7,9 +7,6 @@
 
 from mmselfsup.models.algorithms.mae import MAE
 from mmselfsup.structures import SelfSupDataSample
-from mmselfsup.utils import register_all_modules
-
-register_all_modules()
 
 backbone = dict(type='MAEViT', arch='b', patch_size=16, mask_ratio=0.75)
 neck = dict(