Update pretrained

mmengine/config/collect_meta.py

@@ -1,9 +1,10 @@
+# Copyright (c) OpenMMLab. All rights reserved.
 import os.path as osp
+import re
 from typing import Tuple
 
-from mmengine.utils import check_file_exist
 from mmengine.fileio import load
-import re
+from mmengine.utils import check_file_exist
 
 
 def _get_cfg_meta(package_path: str, cfg_dir: str, cfg_file: str) -> dict:
@@ -40,7 +41,7 @@ def _get_cfg_meta(package_path: str, cfg_dir: str, cfg_file: str) -> dict:
 
 
 def _get_external_cfg_path(package_path: str, cfg_dir: str, cfg_file: str):
-    """ Get relative config path from 'metafile.yml' of external package.
+    """Get relative config path from 'metafile.yml' of external package.
 
     Args:
         package_path (str): Path of external package.
@@ -79,7 +80,7 @@ def _parse_external_cfg_path(rel_cfg_path: str) -> Tuple[str, str]:
     Returns:
         Tuple(str, str): Package name and relative config path.
     """
-    if re.match(r"\w*::\w*/\w*", rel_cfg_path) is None:
+    if re.match(r'\w*::\w*/\w*', rel_cfg_path) is None:
         raise ValueError('`_parse_external_cfg_path` is used for parse '
                          'external package, please specify the package name '
                          'and relative config path, just like '

mmengine/config/config.py

@@ -18,10 +18,10 @@
 from yapf.yapflib.yapf_api import FormatCode
 
 from mmengine.fileio import dump, load
-from mmengine.utils import (check_file_exist, import_modules_from_strings,
-                            check_install_package, get_installed_path)
-from .collect_meta import (_parse_external_cfg_path, _parse_rel_cfg_path,
-                           _get_external_cfg_path, _get_external_cfg_base_path)
+from mmengine.utils import (check_file_exist, check_install_package,
+                            get_installed_path, import_modules_from_strings)
+from .collect_meta import (_get_external_cfg_base_path, _get_external_cfg_path,
+                           _parse_external_cfg_path, _parse_rel_cfg_path)
 
 BASE_KEY = '_base_'
 DELETE_KEY = '_delete_'
@@ -430,8 +430,8 @@ def _file2dict(filename: str,
 
             cfg_dict_list = list()
             cfg_text_list = list()
-            for f in base_filename:
-                cfg_path = Config._get_cfg_path(f, filename)
+            for base_file in base_filename:
+                cfg_path = Config._get_cfg_path(base_file, filename)
                 _cfg_dict, _cfg_text = Config._file2dict(cfg_path)
                 cfg_dict_list.append(_cfg_dict)
                 cfg_text_list.append(_cfg_text)

mmengine/config/download.py

@@ -1,15 +1,18 @@
-import os.path as osp
+# Copyright (c) OpenMMLab. All rights reserved.
 import importlib
+import os.path as osp
 
 import torch.nn as nn
 
+from mmengine.runner import load_checkpoint
+from mmengine.utils import check_install_package, get_installed_path
+from .collect_meta import (_get_cfg_meta, _get_external_cfg_base_path,
+                           _parse_external_cfg_path, _parse_rel_cfg_path)
 from .config import Config
-from mmengine.utils import get_installed_path, check_install_package
-from .collect_meta import (_parse_external_cfg_path, _parse_rel_cfg_path,
-                           _get_cfg_meta, _get_external_cfg_base_path)
 
 
-def get_config(rel_cfg_path: str, suffix='.py') -> Config:
+def get_config(rel_cfg_path: str, suffix='.py', pretrained: bool = False)\
+        -> Config:
     """Get config from external package.
 
     Args:
@@ -18,6 +21,8 @@ def get_config(rel_cfg_path: str, suffix='.py') -> Config:
         suffix (str): Suffix of ``rel_cfg_path``. If rel_cfg_path is a base
             cfg, the `suffix` will be used to get the absolute config path.
             Defaults to '.py'.
+        pretrained (bool): Whether to save pretrained model path. Defaults to
+            False.
 
     Returns:
         Config: A `Config` parsed from external package.
@@ -38,15 +43,19 @@ def get_config(rel_cfg_path: str, suffix='.py') -> Config:
     cfg_meta = _get_cfg_meta(package_path, rel_cfg_dir, rel_cfg_file)
     cfg_path = osp.join(package_path, cfg_meta['Config'])
     cfg = Config.fromfile(cfg_path)
-    assert 'Weights' in cfg_meta, 'Cannot find `Weights` in ' \
-                                  'cfg_file.metafile.yml, please check the ' \
-                                  'metafile'
-    cfg.model_path = cfg_meta['Weights']
+    if pretrained:
+        assert 'Weights' in cfg_meta, 'Cannot find `Weights` in ' \
+                                      'cfg_file.metafile.yml, please check ' \
+                                      'the metafile'
+        cfg.model_path = cfg_meta['Weights']
     return cfg
 
 
-def get_model(rel_cfg_path: str, suffix='.py',
-              build_func_name: str = 'build_model', **kwargs) -> nn.Module:
+def get_model(rel_cfg_path: str,
+              suffix='.py',
+              build_func_name: str = 'build_model',
+              pretrained=False,
+              **kwargs) -> nn.Module:
     """Get built model from external package.
 
     Args:
@@ -55,13 +64,14 @@ def get_model(rel_cfg_path: str, suffix='.py',
         suffix (str): Suffix of ``rel_cfg_path``. If rel_cfg_path is a base
             cfg, the `suffix` will be used to get the absolute config path.
             Defaults to '.py'.
-        build_func_name: Name of model build function. Defaults to
+        build_func_name (str): Name of model build function. Defaults to
             'build_model'
+        pretrained (bool): Whether to load pretrained model. Defaults to False.
 
     Returns:
         nn.Module: Built model.
     """
-    cfg = get_config(rel_cfg_path, suffix)
+    cfg = get_config(rel_cfg_path, suffix, pretrained)
     package = rel_cfg_path.split('::')[0]
 
     try:
@@ -74,4 +84,9 @@ def get_model(rel_cfg_path: str, suffix='.py',
         raise RuntimeError(f'`{build_func_name}` is not defined in '
                            f'`{package}.models`')
     model = build_func(cfg.model, **kwargs)
+    if pretrained:
+        assert hasattr(cfg, 'model_path'), 'Cannot find pretrained model. ' \
+                                           f'Please ensure {rel_cfg_path} ' \
+                                           'is not a base config.'
+        model = load_checkpoint(model, cfg.model_path)
     return model

mmengine/utils/__init__.py

@@ -7,12 +7,12 @@
                    requires_executable, requires_package, slice_list,
                    to_1tuple, to_2tuple, to_3tuple, to_4tuple, to_ntuple,
                    tuple_cast)
+from .package_utils import (call_command, check_install_package,
+                            get_installed_path, is_installed)
 from .parrots_wrapper import TORCH_VERSION
 from .path import (check_file_exist, fopen, is_filepath, mkdir_or_exist,
                    scandir, symlink)
 from .version_utils import digit_version, get_git_hash
-from .package_utils import (is_installed, call_command, get_installed_path,
-                            check_install_package)
 
 __all__ = [
     'is_str', 'iter_cast', 'list_cast', 'tuple_cast', 'is_seq_of',

mmengine/utils/package_utils.py

@@ -1,8 +1,10 @@
-from pkg_resources import get_distribution
-import pkg_resources
-import subprocess
-import os.path as osp
+# Copyright (c) OpenMMLab. All rights reserved.
 import importlib
+import os.path as osp
+import subprocess
+
+import pkg_resources
+from pkg_resources import get_distribution
 
 
 def is_installed(package: str) -> bool:
@@ -61,9 +63,9 @@ def call_command(cmd: list) -> None:
     try:
         subprocess.check_call(cmd)
     except Exception as e:
-        raise highlighted_error(e)  # type: ignore
+        raise e  # type: ignore
 
 
 def check_install_package(package: str):
     if not is_installed(package):
-        call_command(['python', '-m', 'pip', 'install', package])
+        call_command(['python', '-m', 'pip', 'install', package])

tests/test_config/test_collect_meta.py

@@ -1,4 +1,6 @@
+# Copyright (c) OpenMMLab. All rights reserved.
 import os.path
+
 import pytest
 
 from mmengine.config.collect_meta import (_get_external_cfg_base_path,
@@ -49,9 +51,3 @@ def test_parse_rel_cfg_path():
     assert rel_cfg_file == 'cfg_file'
     with pytest.raises(AssertionError):
         _parse_rel_cfg_path('error/cfg_dir/cfg_file')
-
-
-
-
-
-