# Copyright (c) OpenMMLab. All rights reserved. import tempfile from functools import partial import mmcv import numpy as np import pytest import torch from packaging import version from mmocr.core.deployment import (ONNXRuntimeDetector, ONNXRuntimeRecognizer, TensorRTDetector, TensorRTRecognizer) from mmocr.models import build_detector @pytest.mark.skipif(torch.__version__ == 'parrots', reason='skip parrots.') @pytest.mark.skipif( version.parse(torch.__version__) < version.parse('1.4.0'), reason='skip if torch=1.3.x') @pytest.mark.skipif( not torch.cuda.is_available(), reason='skip if on cpu device') def test_detector_wrapper(): try: import onnxruntime as ort # noqa: F401 import tensorrt as trt from mmcv.tensorrt import onnx2trt, save_trt_engine except ImportError: pytest.skip('ONNXRuntime or TensorRT is not available.') cfg = dict( model=dict( type='DBNet', backbone=dict( type='ResNet', depth=18, num_stages=4, out_indices=(0, 1, 2, 3), frozen_stages=-1, norm_cfg=dict(type='BN', requires_grad=True), init_cfg=dict( type='Pretrained', checkpoint='torchvision://resnet18'), norm_eval=False, style='caffe'), neck=dict( type='FPNC', in_channels=[64, 128, 256, 512], lateral_channels=256), bbox_head=dict( type='DBHead', text_repr_type='quad', in_channels=256, loss=dict(type='DBLoss', alpha=5.0, beta=10.0, bbce_loss=True)), train_cfg=None, test_cfg=None)) cfg = mmcv.Config(cfg) pytorch_model = build_detector(cfg.model, None, None) # prepare data inputs = torch.rand(1, 3, 224, 224) img_metas = [{ 'img_shape': [1, 3, 224, 224], 'ori_shape': [1, 3, 224, 224], 'pad_shape': [1, 3, 224, 224], 'filename': None, 'scale_factor': np.array([1, 1, 1, 1]) }] pytorch_model.forward = pytorch_model.forward_dummy with tempfile.TemporaryDirectory() as tmpdirname: onnx_path = f'{tmpdirname}/tmp.onnx' with torch.no_grad(): torch.onnx.export( pytorch_model, inputs, onnx_path, input_names=['input'], output_names=['output'], export_params=True, keep_initializers_as_inputs=False, verbose=False, opset_version=11) # TensorRT part def get_GiB(x: int): """return x GiB.""" return x * (1 << 30) trt_path = onnx_path.replace('.onnx', '.trt') min_shape = [1, 3, 224, 224] max_shape = [1, 3, 224, 224] # create trt engine and wrapper opt_shape_dict = {'input': [min_shape, min_shape, max_shape]} max_workspace_size = get_GiB(1) trt_engine = onnx2trt( onnx_path, opt_shape_dict, log_level=trt.Logger.ERROR, fp16_mode=False, max_workspace_size=max_workspace_size) save_trt_engine(trt_engine, trt_path) print(f'Successfully created TensorRT engine: {trt_path}') wrap_onnx = ONNXRuntimeDetector(onnx_path, cfg, 0) wrap_trt = TensorRTDetector(trt_path, cfg, 0) assert isinstance(wrap_onnx, ONNXRuntimeDetector) assert isinstance(wrap_trt, TensorRTDetector) with torch.no_grad(): onnx_outputs = wrap_onnx.simple_test(inputs, img_metas, rescale=False) trt_outputs = wrap_onnx.simple_test(inputs, img_metas, rescale=False) assert isinstance(onnx_outputs[0], dict) assert isinstance(trt_outputs[0], dict) assert 'boundary_result' in onnx_outputs[0] assert 'boundary_result' in trt_outputs[0] @pytest.mark.skipif(torch.__version__ == 'parrots', reason='skip parrots.') @pytest.mark.skipif( version.parse(torch.__version__) < version.parse('1.4.0'), reason='skip if torch=1.3.x') @pytest.mark.skipif( not torch.cuda.is_available(), reason='skip if on cpu device') def test_recognizer_wrapper(): try: import onnxruntime as ort # noqa: F401 import tensorrt as trt from mmcv.tensorrt import onnx2trt, save_trt_engine except ImportError: pytest.skip('ONNXRuntime or TensorRT is not available.') cfg = dict( label_convertor=dict( type='CTCConvertor', dict_type='DICT36', with_unknown=False, lower=True), model=dict( type='CRNNNet', preprocessor=None, backbone=dict( type='VeryDeepVgg', leaky_relu=False, input_channels=1), encoder=None, decoder=dict(type='CRNNDecoder', in_channels=512, rnn_flag=True), loss=dict(type='CTCLoss'), label_convertor=dict( type='CTCConvertor', dict_type='DICT36', with_unknown=False, lower=True), pretrained=None), train_cfg=None, test_cfg=None) cfg = mmcv.Config(cfg) pytorch_model = build_detector(cfg.model, None, None) # prepare data inputs = torch.rand(1, 1, 32, 32) img_metas = [{ 'img_shape': [1, 1, 32, 32], 'ori_shape': [1, 1, 32, 32], 'pad_shape': [1, 1, 32, 32], 'filename': None, 'scale_factor': np.array([1, 1, 1, 1]) }] pytorch_model.forward = partial( pytorch_model.forward, img_metas=img_metas, return_loss=False, rescale=True) with tempfile.TemporaryDirectory() as tmpdirname: onnx_path = f'{tmpdirname}/tmp.onnx' with torch.no_grad(): torch.onnx.export( pytorch_model, inputs, onnx_path, input_names=['input'], output_names=['output'], export_params=True, keep_initializers_as_inputs=False, verbose=False, opset_version=11) # TensorRT part def get_GiB(x: int): """return x GiB.""" return x * (1 << 30) trt_path = onnx_path.replace('.onnx', '.trt') min_shape = [1, 1, 32, 32] max_shape = [1, 1, 32, 32] # create trt engine and wrapper opt_shape_dict = {'input': [min_shape, min_shape, max_shape]} max_workspace_size = get_GiB(1) trt_engine = onnx2trt( onnx_path, opt_shape_dict, log_level=trt.Logger.ERROR, fp16_mode=False, max_workspace_size=max_workspace_size) save_trt_engine(trt_engine, trt_path) print(f'Successfully created TensorRT engine: {trt_path}') wrap_onnx = ONNXRuntimeRecognizer(onnx_path, cfg, 0) wrap_trt = TensorRTRecognizer(trt_path, cfg, 0) assert isinstance(wrap_onnx, ONNXRuntimeRecognizer) assert isinstance(wrap_trt, TensorRTRecognizer) with torch.no_grad(): onnx_outputs = wrap_onnx.simple_test(inputs, img_metas, rescale=False) trt_outputs = wrap_onnx.simple_test(inputs, img_metas, rescale=False) assert isinstance(onnx_outputs[0], dict) assert isinstance(trt_outputs[0], dict) assert 'text' in onnx_outputs[0] assert 'text' in trt_outputs[0]