Initialize a model 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. Returns: nn.Module: The c
(config, checkpoint=None, device='cuda:0')
| 12 | |
| 13 | |
| 14 | def init_model(config, checkpoint=None, device='cuda:0'): |
| 15 | """Initialize a model from config file. |
| 16 | |
| 17 | Args: |
| 18 | config (str or :obj:`mmcv.Config`): Config file path or the config |
| 19 | object. |
| 20 | checkpoint (str, optional): Checkpoint path. If left as None, the model |
| 21 | will not load any weights. |
| 22 | |
| 23 | Returns: |
| 24 | nn.Module: The constructed model. |
| 25 | (nn.Module, None): The constructed extractor model |
| 26 | """ |
| 27 | if isinstance(config, str): |
| 28 | config = mmcv.Config.fromfile(config) |
| 29 | elif not isinstance(config, mmcv.Config): |
| 30 | raise TypeError('config must be a filename or Config object, ' |
| 31 | f'but got {type(config)}') |
| 32 | config.data.test.test_mode = True |
| 33 | |
| 34 | model = build_architecture(config.model) |
| 35 | if checkpoint is not None: |
| 36 | # load model checkpoint |
| 37 | load_checkpoint(model, checkpoint, map_location=device) |
| 38 | # save the config in the model for convenience |
| 39 | model.cfg = config |
| 40 | model.to(device) |
| 41 | model.eval() |
| 42 | |
| 43 | extractor = None |
| 44 | if config.model.type == 'VideoBodyModelEstimator': |
| 45 | extractor = build_backbone(config.extractor.backbone) |
| 46 | if config.extractor.checkpoint is not None: |
| 47 | # load model checkpoint |
| 48 | load_checkpoint(extractor, config.extractor.checkpoint) |
| 49 | extractor.cfg = config |
| 50 | extractor.to(device) |
| 51 | extractor.eval() |
| 52 | return model, extractor |
| 53 | |
| 54 | |
| 55 | class LoadImage: |
nothing calls this directly
no test coverage detected