Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,963 Bytes
28c256d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import importlib
import os.path as osp
from mmengine.config import Config
from mmengine.config.utils import (_get_cfg_metainfo,
_get_external_cfg_base_path,
_get_package_and_cfg_path)
from mmengine.registry import MODELS, DefaultScope
from mmengine.runner import load_checkpoint
from mmengine.utils import get_installed_path, install_package
def get_config(cfg_path: str, pretrained: bool = False) -> Config:
"""Get config from external package.
Args:
cfg_path (str): External relative config path.
pretrained (bool): Whether to save pretrained model path. If
``pretrained==True``, the url of pretrained model can be accessed
by ``cfg.model_path``. Defaults to False.
Examples:
>>> cfg = get_config('mmdet::faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py', pretrained=True)
>>> # Equivalent to
>>> # cfg = Config.fromfile('/path/to/faster-rcnn_r50_fpn_1x_coco.py')
>>> cfg.model_path
https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth
Returns:
Config: A `Config` parsed from external package.
""" # noqa E301
# Get package name and relative config path.
package, cfg_path = _get_package_and_cfg_path(cfg_path)
# Install package if it's not installed.
install_package(package)
package_path = get_installed_path(package)
try:
# Use `cfg_path` to search target config file.
cfg_meta = _get_cfg_metainfo(package_path, cfg_path)
cfg_path = osp.join(package_path, '.mim', cfg_meta['Config'])
cfg = Config.fromfile(cfg_path)
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']
except ValueError:
# Since the base config does not contain a metafile, the absolute
# config is `osp.join(package_path, cfg_path_prefix, cfg_name)`
cfg_path = _get_external_cfg_base_path(package_path, cfg_path)
cfg = Config.fromfile(cfg_path)
except Exception as e:
raise e
return cfg
def get_model(cfg_path: str, pretrained: bool = False, **kwargs):
"""Get built model from external package.
Args:
cfg_path (str): External relative config path with prefix
'package::' and without suffix.
pretrained (bool): Whether to load pretrained model. Defaults to False.
kwargs (dict): Default arguments to build model.
Examples:
>>> model = get_model('mmdet::faster_rcnn/faster-rcnn_r50_fpn_1x_coco.py', pretrained=True)
>>> type(model)
<class 'mmdet.models.detectors.faster_rcnn.FasterRCNN'>
Returns:
nn.Module: Built model.
""" # noqa E301
package = cfg_path.split('::')[0]
with DefaultScope.overwrite_default_scope(package): # type: ignore
cfg = get_config(cfg_path, pretrained)
if 'data_preprocessor' in cfg:
cfg.model.data_preprocessor = cfg.data_preprocessor
models_module = importlib.import_module(f'{package}.utils')
models_module.register_all_modules() # type: ignore
model = MODELS.build(cfg.model, default_args=kwargs)
if pretrained:
load_checkpoint(model, cfg.model_path)
# Hack to use pretrained weights.
# If we do not set _is_init here, Runner will call
# `model.init_weights()` to overwrite the pretrained model.
model._is_init = True
return model
|