mirrors
/
mmclassification
mirror of https://github.com/open-mmlab/mmclassification.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

support options to override configs in inference.py (#96)

pull/99/head
LXXXXR 2020-11-25 16:14:01 +08:00 committed by GitHub
parent 21fd5019fb
commit c0e7512969
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mmcls/apis/inference.py
View File

@ -11,7 +11,7 @@ from mmcls.datasets.pipelines import Compose
from mmcls.models import build_classifier from mmcls.models import build_classifier
def init_model(config, checkpoint=None, device='cuda:0'): def init_model(config, checkpoint=None, device='cuda:0', options=None):
"""Initialize a classifier from config file. """Initialize a classifier from config file.
Args: Args:
@ -19,6 +19,7 @@ def init_model(config, checkpoint=None, device='cuda:0'):
object. object.
checkpoint (str, optional): Checkpoint path. If left as None, the model checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights. will not load any weights.
options (dict): Options to override some settings in the used config.
Returns: Returns:
nn.Module: The constructed classifier. nn.Module: The constructed classifier.
@ -28,6 +29,8 @@ def init_model(config, checkpoint=None, device='cuda:0'):
elif not isinstance(config, mmcv.Config): elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, ' raise TypeError('config must be a filename or Config object, '
f'but got {type(config)}') f'but got {type(config)}')
if options is not None:
config.merge_from_dict(options)
config.model.pretrained = None config.model.pretrained = None
model = build_classifier(config.model) model = build_classifier(config.model)
if checkpoint is not None: if checkpoint is not None: