|
import argparse |
|
import os |
|
import os.path as osp |
|
|
|
import numpy as np |
|
import onnx |
|
import onnxruntime as ort |
|
import torch |
|
from mmcv.ops import get_onnxruntime_op_path |
|
from mmcv.tensorrt import (TRTWraper, is_tensorrt_plugin_loaded, onnx2trt, |
|
save_trt_engine) |
|
|
|
from mmdet.core import get_classes |
|
from mmdet.core.export import preprocess_example_input |
|
from mmdet.core.visualization.image import imshow_det_bboxes |
|
|
|
|
|
def get_GiB(x: int): |
|
"""return x GiB.""" |
|
return x * (1 << 30) |
|
|
|
|
|
def onnx2tensorrt(onnx_file, |
|
trt_file, |
|
input_config, |
|
verify=False, |
|
show=False, |
|
dataset='coco', |
|
workspace_size=1, |
|
verbose=False): |
|
import tensorrt as trt |
|
onnx_model = onnx.load(onnx_file) |
|
input_shape = input_config['input_shape'] |
|
|
|
opt_shape_dict = {'input': [input_shape, input_shape, input_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=False, |
|
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: |
|
one_img, one_meta = preprocess_example_input(input_config) |
|
input_img_cpu = one_img.detach().cpu().numpy() |
|
input_img_cuda = one_img.cuda() |
|
img = one_meta['show_img'] |
|
|
|
|
|
ort_custom_op_path = get_onnxruntime_op_path() |
|
session_options = ort.SessionOptions() |
|
if osp.exists(ort_custom_op_path): |
|
session_options.register_custom_ops_library(ort_custom_op_path) |
|
sess = ort.InferenceSession(onnx_file, session_options) |
|
output_names = [_.name for _ in sess.get_outputs()] |
|
ort_outputs = sess.run(None, { |
|
'input': input_img_cpu, |
|
}) |
|
with_mask = len(output_names) == 3 |
|
ort_outputs = [_.squeeze(0) for _ in ort_outputs] |
|
ort_dets, ort_labels = ort_outputs[:2] |
|
ort_masks = ort_outputs[2] if with_mask else None |
|
ort_shapes = [_.shape for _ in ort_outputs] |
|
print(f'ONNX Runtime output names: {output_names}, \ |
|
output shapes: {ort_shapes}') |
|
|
|
|
|
trt_model = TRTWraper(trt_file, ['input'], output_names) |
|
with torch.no_grad(): |
|
trt_outputs = trt_model({'input': input_img_cuda}) |
|
trt_outputs = [ |
|
trt_outputs[_].detach().cpu().numpy().squeeze(0) |
|
for _ in output_names |
|
] |
|
trt_dets, trt_labels = trt_outputs[:2] |
|
trt_shapes = [_.shape for _ in trt_outputs] |
|
print(f'TensorRT output names: {output_names}, \ |
|
output shapes: {trt_shapes}') |
|
trt_masks = trt_outputs[2] if with_mask else None |
|
|
|
|
|
if show: |
|
CLASSES = get_classes(dataset) |
|
score_thr = 0.35 |
|
imshow_det_bboxes( |
|
img.copy(), |
|
trt_dets, |
|
trt_labels, |
|
segms=trt_masks, |
|
class_names=CLASSES, |
|
score_thr=score_thr, |
|
win_name='TensorRT') |
|
imshow_det_bboxes( |
|
img.copy(), |
|
ort_dets, |
|
ort_labels, |
|
segms=ort_masks, |
|
class_names=CLASSES, |
|
score_thr=score_thr, |
|
win_name='ONNXRuntime') |
|
|
|
np.testing.assert_allclose(ort_dets, trt_dets, rtol=1e-03, atol=1e-05) |
|
np.testing.assert_allclose(ort_labels, trt_labels) |
|
if with_mask: |
|
np.testing.assert_allclose( |
|
ort_masks, trt_masks, rtol=1e-03, atol=1e-05) |
|
print('The numerical values are the same ' + |
|
'between ONNXRuntime and TensorRT') |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser( |
|
description='Convert MMDetection models from ONNX to TensorRT') |
|
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') |
|
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.') |
|
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') |
|
parser.add_argument( |
|
'--std', |
|
type=float, |
|
nargs='+', |
|
default=[58.395, 57.12, 57.375], |
|
help='Variance value used for preprocess input data') |
|
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() |
|
|
|
if not args.input_img: |
|
args.input_img = osp.join(osp.dirname(__file__), '../demo/demo.jpg') |
|
|
|
if len(args.shape) == 1: |
|
input_shape = (1, 3, args.shape[0], args.shape[0]) |
|
elif len(args.shape) == 2: |
|
input_shape = (1, 3) + tuple(args.shape) |
|
else: |
|
raise ValueError('invalid input shape') |
|
|
|
assert len(args.mean) == 3 |
|
assert len(args.std) == 3 |
|
|
|
normalize_cfg = {'mean': args.mean, 'std': args.std, 'to_rgb': args.to_rgb} |
|
input_config = { |
|
'input_shape': input_shape, |
|
'input_path': args.input_img, |
|
'normalize_cfg': normalize_cfg |
|
} |
|
|
|
|
|
onnx2tensorrt( |
|
args.model, |
|
args.trt_file, |
|
input_config, |
|
verify=args.verify, |
|
show=args.show, |
|
dataset=args.dataset, |
|
workspace_size=args.workspace_size, |
|
verbose=args.verbose) |
|
|
