# Copyright (c) OpenMMLab. All rights reserved. import argparse import os import os.path as osp import warnings import mmcv import torch from mmcv import Config, DictAction from mmcv.cnn import fuse_conv_bn from mmcv.fileio.io import file_handlers from mmcv.runner import get_dist_info, init_dist, load_checkpoint from mmcv.runner.fp16_utils import wrap_fp16_model from mmaction.datasets import build_dataloader, build_dataset from mmaction.models import build_model from mmaction.utils import (build_ddp, build_dp, default_device, register_module_hooks, setup_multi_processes) # TODO import test functions from mmcv and delete them from mmaction2 try: from mmcv.engine import multi_gpu_test, single_gpu_test except (ImportError, ModuleNotFoundError): warnings.warn( 'DeprecationWarning: single_gpu_test, multi_gpu_test, ' 'collect_results_cpu, collect_results_gpu from mmaction2 will be ' 'deprecated. Please install mmcv through master branch.') from mmaction.apis import multi_gpu_test, single_gpu_test def parse_args(): parser = argparse.ArgumentParser( description='MMAction2 test (and eval) a model') parser.add_argument('config', help='test config file path') parser.add_argument('checkpoint', help='checkpoint file') parser.add_argument( '--out', default=None, help='output result file in pkl/yaml/json format') parser.add_argument( '--fuse-conv-bn', action='store_true', help='Whether to fuse conv and bn, this will slightly increase' 'the inference speed') parser.add_argument( '--eval', type=str, nargs='+', help='evaluation metrics, which depends on the dataset, e.g.,' ' "top_k_accuracy", "mean_class_accuracy" for video dataset') parser.add_argument( '--gpu-collect', action='store_true', help='whether to use gpu to collect results') parser.add_argument( '--tmpdir', help='tmp directory used for collecting results from multiple ' 'workers, available when gpu-collect is not specified') parser.add_argument( '--options', nargs='+', action=DictAction, default={}, help='custom options for evaluation, the key-value pair in xxx=yyy ' 'format will be kwargs for dataset.evaluate() function (deprecate), ' 'change to --eval-options instead.') parser.add_argument( '--eval-options', nargs='+', action=DictAction, default={}, help='custom options for evaluation, the key-value pair in xxx=yyy ' 'format will be kwargs for dataset.evaluate() function') parser.add_argument( '--cfg-options', nargs='+', action=DictAction, default={}, help='override some settings in the used config, the key-value pair ' 'in xxx=yyy format will be merged into config file. For example, ' "'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'") parser.add_argument( '--average-clips', choices=['score', 'prob', None], default=None, help='average type when averaging test clips') parser.add_argument( '--launcher', choices=['none', 'pytorch', 'slurm', 'mpi'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) parser.add_argument( '--onnx', action='store_true', help='Whether to test with onnx model or not') parser.add_argument( '--tensorrt', action='store_true', help='Whether to test with TensorRT engine or not') args = parser.parse_args() if 'LOCAL_RANK' not in os.environ: os.environ['LOCAL_RANK'] = str(args.local_rank) if args.options and args.eval_options: raise ValueError( '--options and --eval-options cannot be both ' 'specified, --options is deprecated in favor of --eval-options') if args.options: warnings.warn('--options is deprecated in favor of --eval-options') args.eval_options = args.options return args def turn_off_pretrained(cfg): # recursively find all pretrained in the model config, # and set them None to avoid redundant pretrain steps for testing if 'pretrained' in cfg: cfg.pretrained = None # recursively turn off pretrained value for sub_cfg in cfg.values(): if isinstance(sub_cfg, dict): turn_off_pretrained(sub_cfg) def inference_pytorch(args, cfg, distributed, data_loader): """Get predictions by pytorch models.""" if args.average_clips is not None: # You can set average_clips during testing, it will override the # original setting if cfg.model.get('test_cfg') is None and cfg.get('test_cfg') is None: cfg.model.setdefault('test_cfg', dict(average_clips=args.average_clips)) else: if cfg.model.get('test_cfg') is not None: cfg.model.test_cfg.average_clips = args.average_clips else: cfg.test_cfg.average_clips = args.average_clips # remove redundant pretrain steps for testing turn_off_pretrained(cfg.model) # build the model and load checkpoint model = build_model( cfg.model, train_cfg=None, test_cfg=cfg.get('test_cfg')) if len(cfg.module_hooks) > 0: register_module_hooks(model, cfg.module_hooks) fp16_cfg = cfg.get('fp16', None) if fp16_cfg is not None: wrap_fp16_model(model) load_checkpoint(model, args.checkpoint, map_location='cpu') if args.fuse_conv_bn: model = fuse_conv_bn(model) if not distributed: model = build_dp( model, default_device, default_args=dict(device_ids=cfg.gpu_ids)) outputs = single_gpu_test(model, data_loader) else: model = build_ddp( model, default_device, default_args=dict( device_ids=[int(os.environ['LOCAL_RANK'])], broadcast_buffers=False)) outputs = multi_gpu_test(model, data_loader, args.tmpdir, args.gpu_collect) return outputs def inference_tensorrt(ckpt_path, distributed, data_loader, batch_size): """Get predictions by TensorRT engine. For now, multi-gpu mode and dynamic tensor shape are not supported. """ assert not distributed, \ 'TensorRT engine inference only supports single gpu mode.' import tensorrt as trt from mmcv.tensorrt.tensorrt_utils import (torch_device_from_trt, torch_dtype_from_trt) # load engine with trt.Logger() as logger, trt.Runtime(logger) as runtime: with open(ckpt_path, mode='rb') as f: engine_bytes = f.read() engine = runtime.deserialize_cuda_engine(engine_bytes) # For now, only support fixed input tensor cur_batch_size = engine.get_binding_shape(0)[0] assert batch_size == cur_batch_size, \ ('Dataset and TensorRT model should share the same batch size, ' f'but get {batch_size} and {cur_batch_size}') context = engine.create_execution_context() # get output tensor dtype = torch_dtype_from_trt(engine.get_binding_dtype(1)) shape = tuple(context.get_binding_shape(1)) device = torch_device_from_trt(engine.get_location(1)) output = torch.empty( size=shape, dtype=dtype, device=device, requires_grad=False) # get predictions results = [] dataset = data_loader.dataset prog_bar = mmcv.ProgressBar(len(dataset)) for data in data_loader: bindings = [ data['imgs'].contiguous().data_ptr(), output.contiguous().data_ptr() ] context.execute_async_v2(bindings, torch.cuda.current_stream().cuda_stream) results.extend(output.cpu().numpy()) batch_size = len(next(iter(data.values()))) for _ in range(batch_size): prog_bar.update() return results def inference_onnx(ckpt_path, distributed, data_loader, batch_size): """Get predictions by ONNX. For now, multi-gpu mode and dynamic tensor shape are not supported. """ assert not distributed, 'ONNX inference only supports single gpu mode.' import onnx import onnxruntime as rt # get input tensor name onnx_model = onnx.load(ckpt_path) input_all = [node.name for node in onnx_model.graph.input] input_initializer = [node.name for node in onnx_model.graph.initializer] net_feed_input = list(set(input_all) - set(input_initializer)) assert len(net_feed_input) == 1 # For now, only support fixed tensor shape input_tensor = None for tensor in onnx_model.graph.input: if tensor.name == net_feed_input[0]: input_tensor = tensor break cur_batch_size = input_tensor.type.tensor_type.shape.dim[0].dim_value assert batch_size == cur_batch_size, \ ('Dataset and ONNX model should share the same batch size, ' f'but get {batch_size} and {cur_batch_size}') # get predictions sess = rt.InferenceSession(ckpt_path) results = [] dataset = data_loader.dataset prog_bar = mmcv.ProgressBar(len(dataset)) for data in data_loader: imgs = data['imgs'].cpu().numpy() onnx_result = sess.run(None, {net_feed_input[0]: imgs})[0] results.extend(onnx_result) batch_size = len(next(iter(data.values()))) for _ in range(batch_size): prog_bar.update() return results def main(): args = parse_args() if args.tensorrt and args.onnx: raise ValueError( 'Cannot set onnx mode and tensorrt mode at the same time.') cfg = Config.fromfile(args.config) cfg.merge_from_dict(args.cfg_options) # set multi-process settings setup_multi_processes(cfg) # Load output_config from cfg output_config = cfg.get('output_config', {}) if args.out: # Overwrite output_config from args.out output_config = Config._merge_a_into_b( dict(out=args.out), output_config) # Load eval_config from cfg eval_config = cfg.get('eval_config', {}) if args.eval: # Overwrite eval_config from args.eval eval_config = Config._merge_a_into_b( dict(metrics=args.eval), eval_config) if args.eval_options: # Add options from args.eval_options eval_config = Config._merge_a_into_b(args.eval_options, eval_config) assert output_config or eval_config, \ ('Please specify at least one operation (save or eval the ' 'results) with the argument "--out" or "--eval"') dataset_type = cfg.data.test.type if output_config.get('out', None): if 'output_format' in output_config: # ugly workround to make recognition and localization the same warnings.warn( 'Skip checking `output_format` in localization task.') else: out = output_config['out'] # make sure the dirname of the output path exists mmcv.mkdir_or_exist(osp.dirname(out)) _, suffix = osp.splitext(out) if dataset_type == 'AVADataset': assert suffix[1:] == 'csv', ('For AVADataset, the format of ' 'the output file should be csv') else: assert suffix[1:] in file_handlers, ( 'The format of the output ' 'file should be json, pickle or yaml') # set cudnn benchmark if cfg.get('cudnn_benchmark', False): torch.backends.cudnn.benchmark = True cfg.data.test.test_mode = True # init distributed env first, since logger depends on the dist info. if args.launcher == 'none': distributed = False else: distributed = True init_dist(args.launcher, **cfg.dist_params) # The flag is used to register module's hooks cfg.setdefault('module_hooks', []) # build the dataloader dataset = build_dataset(cfg.data.test, dict(test_mode=True)) dataloader_setting = dict( videos_per_gpu=cfg.data.get('videos_per_gpu', 1), workers_per_gpu=cfg.data.get('workers_per_gpu', 1), dist=distributed, shuffle=False) dataloader_setting = dict(dataloader_setting, **cfg.data.get('test_dataloader', {})) data_loader = build_dataloader(dataset, **dataloader_setting) if args.tensorrt: outputs = inference_tensorrt(args.checkpoint, distributed, data_loader, dataloader_setting['videos_per_gpu']) elif args.onnx: outputs = inference_onnx(args.checkpoint, distributed, data_loader, dataloader_setting['videos_per_gpu']) else: outputs = inference_pytorch(args, cfg, distributed, data_loader) rank, _ = get_dist_info() if rank == 0: if output_config.get('out', None): out = output_config['out'] print(f'\nwriting results to {out}') dataset.dump_results(outputs, **output_config) if eval_config: eval_res = dataset.evaluate(outputs, **eval_config) for name, val in eval_res.items(): print(f'{name}: {val:.04f}') if __name__ == '__main__': main()