Python runner.load_checkpoint方法代码示例

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def init_model(config, checkpoint=None, device='cuda:0'):
    """
    Initialize a stereo 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 constructed stereo model.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        'but got {}'.format(type(config)))

    model = build_model(config)
    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint)
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()
    return model 

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def init_weights(self, pretrained=None):
        print("init hrnet weights")
#         if isinstance(pretrained, str):
#             logger = logging.getLogger()
#             load_checkpoint(self, pretrained, strict=False, logger=logger)
#         elif pretrained is None:
#             for m in self.modules():
#                 if isinstance(m, nn.Conv2d):
#                     kaiming_init(m)
#                 elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
#                     constant_init(m, 1)

#             if self.zero_init_residual:
#                 for m in self.modules():
#                     if isinstance(m, Bottleneck):
#                         constant_init(m.norm3, 0)
#                     elif isinstance(m, BasicBlock):
#                         constant_init(m.norm2, 0)
#         else:
#             raise TypeError('pretrained must be a str or None') 

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def init_detector(config, checkpoint=None, device='cuda:0'):
    """Initialize a detector 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 constructed detector.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {type(config)}')
    config.model.pretrained = None
    model = build_detector(config.model, test_cfg=config.test_cfg)
    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint)
        if 'CLASSES' in checkpoint['meta']:
            model.CLASSES = checkpoint['meta']['CLASSES']
        else:
            warnings.simplefilter('once')
            warnings.warn('Class names are not saved in the checkpoint\'s '
                          'meta data, use COCO classes by default.')
            model.CLASSES = get_classes('coco')
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()
    return model 

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def init_detector(config, checkpoint=None, device='cuda:0'):
    """Initialize a detector 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 constructed detector.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        'but got {}'.format(type(config)))
    config.model.pretrained = None
    model = build_detector(config.model, test_cfg=config.test_cfg)
    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint)
        if 'CLASSES' in checkpoint['meta']:
            model.CLASSES = checkpoint['meta']['CLASSES']
        else:
            warnings.warn('Class names are not saved in the checkpoint\'s '
                          'meta data, use COCO classes by default.')
            model.CLASSES = get_classes('coco')
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()
    return model 

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def init_detector(config, checkpoint=None, device='cuda:0'):
    """Initialize a detector 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 constructed detector.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        'but got {}'.format(type(config)))
    config.model.pretrained = None
    model = build_detector(config.model, test_cfg=config.test_cfg)
    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint)
        if 'CLASSES' in checkpoint['meta']:
            model.CLASSES = checkpoint['meta']['classes']
        else:
            warnings.warn('Class names are not saved in the checkpoint\'s '
                          'meta data, use COCO classes by default.')
            model.CLASSES = get_classes('coco')
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()
    return model 

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def load_checkpoint(model, filename, *args, **kwargs):
    try:
        filename = get_mmskeleton_url(filename)
        return mmcv_load_checkpoint(model, filename, *args, **kwargs)
    except (urllib.error.HTTPError, urllib.error.URLError) as e:
        raise Exception(url_error_message.format(filename)) from e 

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def load_model():
    model = build_detector(cfg.model, test_cfg=cfg.test_cfg)
    _ = load_checkpoint(model, model_cfgs[0][1]) # 7 it/s
    return model 

# 需要导入模块: from mmcv import runner [as 别名]
# 或者: from mmcv.runner import load_checkpoint [as 别名]
def _make_stage(self, layer_config, in_channels, multiscale_output=True):
        num_modules = layer_config['num_modules']
        num_branches = layer_config['num_branches']
        num_blocks = layer_config['num_blocks']
        num_channels = layer_config['num_channels']
        block = self.blocks_dict[layer_config['block']]

        hr_modules = []
        for i in range(num_modules):
            # multi_scale_output is only used for the last module
            if not multiscale_output and i == num_modules - 1:
                reset_multiscale_output = False
            else:
                reset_multiscale_output = True

            hr_modules.append(
                HRModule(
                    num_branches,
                    block,
                    num_blocks,
                    in_channels,
                    num_channels,
                    reset_multiscale_output,
                    with_cp=self.with_cp,
                    norm_cfg=self.norm_cfg,
                    conv_cfg=self.conv_cfg))

        return nn.Sequential(*hr_modules), in_channels

#     def init_weights(self, pretrained=None):
#         if isinstance(pretrained, str):
#             logger = logging.getLogger()
#             load_checkpoint(self, pretrained, strict=False, logger=logger)
#         elif pretrained is None:
#             for m in self.modules():
#                 if isinstance(m, nn.Conv2d):
#                     kaiming_init(m)
#                 elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
#                     constant_init(m, 1)

#             if self.zero_init_residual:
#                 for m in self.modules():
#                     if isinstance(m, Bottleneck):
#                         constant_init(m.norm3, 0)
#                     elif isinstance(m, BasicBlock):
#                         constant_init(m.norm2, 0)
#         else:
#             raise TypeError('pretrained must be a str or None')