# Copyright (c) Alibaba, Inc. and its affiliates. """ isort:skip_file """ import argparse import os import os.path as osp import sys import time import requests import torch try: from blade_compression.fx_quantization.prepare import (convert, enable_calibration, prepare_fx) except ImportError: raise ImportError( 'Please read docs and run "pip install http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/third_party/blade_compression-0.0.1-py3-none-any.whl" ' 'to install blade_compression') from mmcv.parallel import MMDataParallel from mmcv.runner import get_dist_info from easycv.models import build_model from easycv.apis import single_cpu_test, single_gpu_test from easycv.core.evaluation.builder import build_evaluator from easycv.datasets import build_dataloader, build_dataset from easycv.utils.logger import get_root_logger from easycv.utils.flops_counter import get_model_info from easycv.file import io from easycv.toolkit.quantize.quantize_utils import calib, quantize_config_check from easycv.utils.checkpoint import load_checkpoint from easycv.utils.config_tools import (CONFIG_TEMPLATE_ZOO, mmcv_config_fromfile, rebuild_config) sys.path.append(os.path.abspath(os.path.dirname(os.path.dirname(__file__)))) sys.path.append( os.path.abspath( osp.join(os.path.dirname(os.path.dirname(__file__)), '../'))) def parse_args(): parser = argparse.ArgumentParser(description='EasyCV quantize a model') parser.add_argument('config', help='model config file path') parser.add_argument('checkpoint', help='checkpoint file') parser.add_argument( '--work_dir', type=str, default=None, help='the dir to save quantized models') parser.add_argument( '--model_type', type=str, default=None, help= 'parameterize param when user specific choose a model config template like CLASSIFICATION: classification.py' ) parser.add_argument('--local_rank', type=int, default=0) parser.add_argument( '--device', type=str, default='cpu', help='the device quantized models use') parser.add_argument( '--backend', type=str, default='PyTorch', help="the quantized models's framework") parser.add_argument( '--user_config_params', nargs=argparse.REMAINDER, default=None, help='modify config options using the command-line') args = parser.parse_args() return args def quantize_eval(cfg, model, eval_mode): for eval_pipe in cfg.eval_pipelines: eval_data = eval_pipe.data # build the dataloader 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) if eval_mode == 'cuda': outputs = single_gpu_test(model, data_loader, mode=eval_pipe.mode) elif eval_mode == 'cpu': outputs = single_cpu_test(model, data_loader, mode=eval_pipe.mode) rank, _ = get_dist_info() if rank == 0: for t in eval_pipe.evaluators: if 'metric_type' in t: t.pop('metric_type') evaluators = build_evaluator(eval_pipe.evaluators) eval_result = dataset.evaluate(outputs, evaluators=evaluators) print(f'\n eval_result {eval_result}') 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) if __name__ == '__main__': main()