|
|
|
|
|
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 |
|
|
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) |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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] |
|
|
|
|
|
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 |
|
|
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) |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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] |
|
|
|
|
|
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] |
|
|
|
