mmdetection-2.26.0/tools/deployment/onnx2tensorrt.py

# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os
import os.path as osp
import warnings

import numpy as np
import onnx
import torch
from mmcv import Config
from mmcv.tensorrt import is_tensorrt_plugin_loaded, onnx2trt, save_trt_engine

from mmdet.core.export import preprocess_example_input
from mmdet.core.export.model_wrappers import (ONNXRuntimeDetector,
 TensorRTDetector)
from mmdet.datasets import DATASETS


def get_GiB(x: int):
 """return x GiB."""
 return x * (1 << 30)


def onnx2tensorrt(onnx_file,
 trt_file,
 input_config,
 verify=False,
 show=False,
 workspace_size=1,
 verbose=False):
 import tensorrt as trt
 onnx_model = onnx.load(onnx_file)
 max_shape = input_config['max_shape']
 min_shape = input_config['min_shape']
 opt_shape = input_config['opt_shape']
 fp16_mode = False
 # create trt engine and wrapper
 opt_shape_dict = {'input': [min_shape, opt_shape, max_shape]}
 max_workspace_size = get_GiB(workspace_size)
 trt_engine = onnx2trt(
 onnx_model,
 opt_shape_dict,
 log_level=trt.Logger.VERBOSE if verbose else trt.Logger.ERROR,
 fp16_mode=fp16_mode,
 max_workspace_size=max_workspace_size)
 save_dir, _ = osp.split(trt_file)
 if save_dir:
 os.makedirs(save_dir, exist_ok=True)
 save_trt_engine(trt_engine, trt_file)
 print(f'Successfully created TensorRT engine: {trt_file}')

 if verify:
 # prepare input
 one_img, one_meta = preprocess_example_input(input_config)
 img_list, img_meta_list = [one_img], [[one_meta]]
 img_list = [_.cuda().contiguous() for _ in img_list]

 # wrap ONNX and TensorRT model
 onnx_model = ONNXRuntimeDetector(onnx_file, CLASSES, device_id=0)
 trt_model = TensorRTDetector(trt_file, CLASSES, device_id=0)

 # inference with wrapped model
 with torch.no_grad():
 onnx_results = onnx_model(
 img_list, img_metas=img_meta_list, return_loss=False)[0]
 trt_results = trt_model(
 img_list, img_metas=img_meta_list, return_loss=False)[0]

 if show:
 out_file_ort, out_file_trt = None, None
 else:
 out_file_ort, out_file_trt = 'show-ort.png', 'show-trt.png'
 show_img = one_meta['show_img']
 score_thr = 0.3
 onnx_model.show_result(
 show_img,
 onnx_results,
 score_thr=score_thr,
 show=True,
 win_name='ONNXRuntime',
 out_file=out_file_ort)
 trt_model.show_result(
 show_img,
 trt_results,
 score_thr=score_thr,
 show=True,
 win_name='TensorRT',
 out_file=out_file_trt)
 with_mask = trt_model.with_masks
 # compare a part of result
 if with_mask:
 compare_pairs = list(zip(onnx_results, trt_results))
 else:
 compare_pairs = [(onnx_results, trt_results)]
 err_msg = 'The numerical values are different between Pytorch' + \
 ' and ONNX, but it does not necessarily mean the' + \
 ' exported ONNX model is problematic.'
 # check the numerical value
 for onnx_res, pytorch_res in compare_pairs:
 for o_res, p_res in zip(onnx_res, pytorch_res):
 np.testing.assert_allclose(
 o_res, p_res, rtol=1e-03, atol=1e-05, err_msg=err_msg)
 print('The numerical values are the same between Pytorch and ONNX')


def parse_normalize_cfg(test_pipeline):
 transforms = None
 for pipeline in test_pipeline:
 if 'transforms' in pipeline:
 transforms = pipeline['transforms']
 break
 assert transforms is not None, 'Failed to find `transforms`'
 norm_config_li = [_ for _ in transforms if _['type'] == 'Normalize']
 assert len(norm_config_li) == 1, '`norm_config` should only have one'
 norm_config = norm_config_li[0]
 return norm_config


def parse_args():
 parser = argparse.ArgumentParser(
 description='Convert MMDetection models from ONNX to TensorRT')
 parser.add_argument('config', help='test config file path')
 parser.add_argument('model', help='Filename of input ONNX model')
 parser.add_argument(
 '--trt-file',
 type=str,
 default='tmp.trt',
 help='Filename of output TensorRT engine')
 parser.add_argument(
 '--input-img', type=str, default='', help='Image for test')
 parser.add_argument(
 '--show', action='store_true', help='Whether to show output results')
 parser.add_argument(
 '--dataset',
 type=str,
 default='coco',
 help='Dataset name. This argument is deprecated and will be \
 removed in future releases.')
 parser.add_argument(
 '--verify',
 action='store_true',
 help='Verify the outputs of ONNXRuntime and TensorRT')
 parser.add_argument(
 '--verbose',
 action='store_true',
 help='Whether to verbose logging messages while creating \
 TensorRT engine. Defaults to False.')
 parser.add_argument(
 '--to-rgb',
 action='store_false',
 help='Feed model with RGB or BGR image. Default is RGB. This \
 argument is deprecated and will be removed in future releases.')
 parser.add_argument(
 '--shape',
 type=int,
 nargs='+',
 default=[400, 600],
 help='Input size of the model')
 parser.add_argument(
 '--mean',
 type=float,
 nargs='+',
 default=[123.675, 116.28, 103.53],
 help='Mean value used for preprocess input data. This argument \
 is deprecated and will be removed in future releases.')
 parser.add_argument(
 '--std',
 type=float,
 nargs='+',
 default=[58.395, 57.12, 57.375],
 help='Variance value used for preprocess input data. \
 This argument is deprecated and will be removed in future releases.')
 parser.add_argument(
 '--min-shape',
 type=int,
 nargs='+',
 default=None,
 help='Minimum input size of the model in TensorRT')
 parser.add_argument(
 '--max-shape',
 type=int,
 nargs='+',
 default=None,
 help='Maximum input size of the model in TensorRT')
 parser.add_argument(
 '--workspace-size',
 type=int,
 default=1,
 help='Max workspace size in GiB')

 args = parser.parse_args()
 return args


if __name__ == '__main__':

 assert is_tensorrt_plugin_loaded(), 'TensorRT plugin should be compiled.'
 args = parse_args()
 warnings.warn(
 'Arguments like `--to-rgb`, `--mean`, `--std`, `--dataset` would be \
 parsed directly from config file and are deprecated and will be \
 removed in future releases.')
 if not args.input_img:
 args.input_img = osp.join(osp.dirname(__file__), '../../demo/demo.jpg')

 cfg = Config.fromfile(args.config)

 def parse_shape(shape):
 if len(shape) == 1:
 shape = (1, 3, shape[0], shape[0])
 elif len(args.shape) == 2:
 shape = (1, 3) + tuple(shape)
 else:
 raise ValueError('invalid input shape')
 return shape

 if args.shape:
 input_shape = parse_shape(args.shape)
 else:
 img_scale = cfg.test_pipeline[1]['img_scale']
 input_shape = (1, 3, img_scale[1], img_scale[0])

 if not args.max_shape:
 max_shape = input_shape
 else:
 max_shape = parse_shape(args.max_shape)

 if not args.min_shape:
 min_shape = input_shape
 else:
 min_shape = parse_shape(args.min_shape)

 dataset = DATASETS.get(cfg.data.test['type'])
 assert (dataset is not None)
 CLASSES = dataset.CLASSES
 normalize_cfg = parse_normalize_cfg(cfg.test_pipeline)

 input_config = {
 'min_shape': min_shape,
 'opt_shape': input_shape,
 'max_shape': max_shape,
 'input_shape': input_shape,
 'input_path': args.input_img,
 'normalize_cfg': normalize_cfg
 }
 # Create TensorRT engine
 onnx2tensorrt(
 args.model,
 args.trt_file,
 input_config,
 verify=args.verify,
 show=args.show,
 workspace_size=args.workspace_size,
 verbose=args.verbose)

 # Following strings of text style are from colorama package
 bright_style, reset_style = '\x1b[1m', '\x1b[0m'
 red_text, blue_text = '\x1b[31m', '\x1b[34m'
 white_background = '\x1b[107m'

 msg = white_background + bright_style + red_text
 msg += 'DeprecationWarning: This tool will be deprecated in future. '
 msg += blue_text + 'Welcome to use the unified model deployment toolbox '
 msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
 msg += reset_style
 warnings.warn(msg)