Skip to content

Commit

Permalink
support options to override configs in inference.py (open-mmlab#96)
Browse files Browse the repository at this point in the history
  • Loading branch information
LXXXXR authored Nov 25, 2020
1 parent 21fd501 commit c0e7512
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion mmcls/apis/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,15 @@
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.
Args:
config (str or :obj:`mmcv.Config`): Config file path or the config
object.
checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights.
options (dict): Options to override some settings in the used config.
Returns:
nn.Module: The constructed classifier.
Expand All @@ -28,6 +29,8 @@ def init_model(config, checkpoint=None, device='cuda:0'):
elif not isinstance(config, mmcv.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)
config.model.pretrained = None
model = build_classifier(config.model)
if checkpoint is not None:
Expand Down

0 comments on commit c0e7512

Please sign in to comment.