def export()

in easycv/apis/export.py [0:0]

• 38 lines of code
• 17 McCabe index (conditional complexity)

def export(cfg, ckpt_path, filename, model=None, **kwargs):
    """ export model for inference

    Args:
        cfg: Config object
        ckpt_path (str): path to checkpoint file
        filename (str): filename to save exported models
        model (nn.module): model instance
    """
    if hasattr(cfg.model, 'pretrained'):
        logging.warning(
            'Export needs to set model.pretrained to false to avoid hanging during distributed training'
        )
        cfg.model.pretrained = False

    if model is None:
        model = build_model(cfg.model)

    if ckpt_path != 'dummy':
        load_checkpoint(model, ckpt_path, map_location='cpu')
    else:
        if hasattr(cfg.model, 'backbone') and hasattr(cfg.model.backbone,
                                                      'pretrained'):
            logging.warning(
                'Export needs to set model.backbone.pretrained to false to avoid hanging during distributed training'
            )
            cfg.model.backbone.pretrained = False

    if isinstance(model, MOCO) or isinstance(model, DINO):
        _export_moco(model, cfg, filename, **kwargs)
    elif isinstance(model, MoBY):
        _export_moby(model, cfg, filename, **kwargs)
    elif isinstance(model, SWAV):
        _export_swav(model, cfg, filename, **kwargs)
    elif isinstance(model, Classification):
        _export_cls(model, cfg, filename, **kwargs)
    elif isinstance(model, YOLOX):
        _export_yolox(model, cfg, filename, **kwargs)
    elif isinstance(model, BEVFormer):
        _export_bevformer(model, cfg, filename, **kwargs)
    elif isinstance(model, TopDown):
        _export_pose_topdown(model, cfg, filename, **kwargs)
    elif isinstance(model, SkeletonGCN):
        _export_stgcn(model, cfg, filename, **kwargs)
    elif hasattr(cfg, 'export') and getattr(cfg.export, 'use_jit', False):
        export_jit_model(model, cfg, filename, **kwargs)
        return
    else:
        _export_common(model, cfg, filename, **kwargs)