def main()

in tools/quantize.py [0:0]

• 101 lines of code
• 29 McCabe index (conditional complexity)

def main():
    args = parse_args()

    if args.model_type is not None and args.config is None:
        assert args.model_type in CONFIG_TEMPLATE_ZOO, 'model_type must be in [%s]' % (
            ', '.join(CONFIG_TEMPLATE_ZOO.keys()))
        print('model_type=%s, config file will be replaced by %s' %
              (args.model_type, CONFIG_TEMPLATE_ZOO[args.model_type]))
        args.config = CONFIG_TEMPLATE_ZOO[args.model_type]

    if args.config.startswith('http'):

        r = requests.get(args.config)
        # download config in current dir
        tpath = args.config.split('/')[-1]
        while not osp.exists(tpath):
            try:
                with open(tpath, 'wb') as code:
                    code.write(r.content)
            except:
                pass

        args.config = tpath

    cfg = mmcv_config_fromfile(args.config)

    if args.user_config_params is not None:
        assert args.model_type is not None, 'model_type must be setted'
        # rebuild config by user config params
        cfg = rebuild_config(cfg, args.user_config_params)

    # check oss_config and init oss io
    if cfg.get('oss_io_config', None) is not None:
        io.access_oss(**cfg.oss_io_config)

    # set cudnn_benchmark
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True

    # update configs according to CLI args
    if args.work_dir is not None:
        cfg.work_dir = args.work_dir

    # if `work_dir` is oss path, redirect `work_dir` to local path, add `oss_work_dir` point to oss path,
    # and use osssync hook to upload log and ckpt in work_dir to oss_work_dir
    if cfg.work_dir.startswith('oss://'):
        cfg.oss_work_dir = cfg.work_dir
        cfg.work_dir = osp.join('work_dirs',
                                cfg.work_dir.replace('oss://', ''))
    else:
        cfg.oss_work_dir = None

    # create work_dir
    if not io.exists(cfg.work_dir):
        io.makedirs(cfg.work_dir)

    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = osp.join(cfg.work_dir, '{}.log'.format(timestamp))
    logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)

    cfg.model.pretrained = None
    if cfg.model.get('neck'):
        if type(cfg.model.neck) is list:
            pass
        else:
            if cfg.model.neck.get('rfp_backbone'):
                if cfg.model.neck.rfp_backbone.get('pretrained'):
                    cfg.model.neck.rfp_backbone.pretrained = None

    # build the model and load checkpoint
    model = build_model(cfg.model)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    logger.info(f'use device {device}')
    checkpoint = load_checkpoint(model, args.checkpoint, map_location=device)
    model.eval()
    model.to(device)
    # old versions did not save class info in checkpoints, this walkaround is
    # for backward compatibility
    if 'meta' in checkpoint and 'CLASSES' in checkpoint['meta']:
        model.CLASSES = checkpoint['meta']['CLASSES']
    elif hasattr(cfg, 'CLASSES'):
        model.CLASSES = cfg.CLASSES

    # MMDataParallel for gpu
    if device == 'cuda':
        base_model = MMDataParallel(model, device_ids=[0])
    else:
        base_model = model

    # eval base model before quantizing
    get_model_info(model, cfg.img_scale, cfg.model, logger)
    quantize_eval(cfg, base_model, device)

    # setting quantize config
    quantize_config = quantize_config_check(args.device, args.backend,
                                            args.model_type)

    model.to('cuda')
    prepared_backbone = prepare_fx(model.backbone.eval(), quantize_config)
    enable_calibration(prepared_backbone)
    # build calib dataloader, only need 50 samples
    logger.info('build calib dataloader')
    eval_data = cfg.eval_pipelines[0].data
    imgs_per_gpu = eval_data.pop('imgs_per_gpu', cfg.data.imgs_per_gpu)

    dataset = build_dataset(eval_data)
    data_loader = build_dataloader(
        dataset,
        imgs_per_gpu=imgs_per_gpu,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=False,
        shuffle=False)

    # guarantee accuracy
    logger.info('guarantee calib')
    calib(prepared_backbone, data_loader)

    # quantized model on cpu
    model.to('cpu')

    # quantizing model
    logger.info('convert model')
    quantized_backbone, _ = convert(prepared_backbone, quantize_config)
    model.backbone = quantized_backbone
    model.eval()

    # cpu eval
    logger.info('quantized model eval')
    get_model_info(model, cfg.img_scale, cfg.model, logger)
    quantize_eval(cfg, model, 'cpu')

    input_shape = (1, 3, cfg.img_scale[0], cfg.img_scale[1])
    model.head.decode_in_inference = False
    dummy = torch.randn(input_shape)
    traced_model = torch.jit.trace(model, dummy)
    model_path = osp.join(cfg.work_dir, 'quantize_model.pt')
    torch.jit.save(traced_model, model_path)

    if cfg.oss_work_dir is not None:
        export_oss_path = os.path.join(cfg.oss_work_dir, 'quantize_model.pt')
        if not os.path.exists(model_path):
            logger.warning(f'{model_path} does not exists, skip upload')
        else:
            logger.info(f'upload {model_path} to {export_oss_path}')
            io.safe_copy(model_path, export_oss_path)