mmdeploy/tests/test_backend/test_wrapper.py
q.yao b32fc41bed
[Refactor][API2.0] Api refactor2.0 (#529)
* [refactor][API2.0]  Add onnx export and jit trace (#419)

* first commit

* add async call

* add new api onnx export and jit trace

* add decorator

* fix ci

* fix torchscript ci

* fix loader

* better pipemanager

* remove comment, better import

* add kwargs

* remove comment

* better pipeline manager

* remove print

* [Refactor][API2.0] Api partition calibration (#433)

* first commit

* add async call

* add new api onnx export and jit trace

* add decorator

* fix ci

* fix torchscript ci

* fix loader

* better pipemanager

* remove comment, better import

* add partition

* move calibration

* Better create_calib_table

* better deploy

* add kwargs

* remove comment

* better pipeline manager

* rename api, remove reduant variable, and misc

* [Refactor][API2.0] Api ncnn openvino (#435)

* first commit

* add async call

* add new api onnx export and jit trace

* add decorator

* fix ci

* fix torchscript ci

* fix loader

* better pipemanager

* remove comment, better import

* add ncnn api

* finish ncnn api

* add openvino support

* add kwargs

* remove comment

* better pipeline manager

* merge fix

* merge util and onnx2ncnn

* fix docstring

* [Refactor][API2.0] API for TensorRT (#519)

* first commit

* add async call

* add new api onnx export and jit trace

* add decorator

* fix ci

* fix torchscript ci

* fix loader

* better pipemanager

* remove comment, better import

* add partition

* move calibration

* Better create_calib_table

* better deploy

* add kwargs

* remove comment

* Add tensorrt API

* better pipeline manager

* add tensorrt new api

* remove print

* rename api, remove reduant variable, and misc

* add docstring

* [Refactor][API2.0] Api ppl other (#528)

* first commit

* add async call

* add new api onnx export and jit trace

* add decorator

* fix ci

* fix torchscript ci

* fix loader

* better pipemanager

* remove comment, better import

* add kwargs

* Add new APIS for pplnn sdk and misc

* remove comment

* better pipeline manager

* merge fix

* update tools/onnx2pplnn.py

* rename function
2022-05-31 09:18:18 +08:00

188 lines
6.3 KiB
Python

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import subprocess
import tempfile
import pytest
import torch
import torch.nn as nn
from mmdeploy.utils.constants import Backend
from mmdeploy.utils.test import check_backend
onnx_file = tempfile.NamedTemporaryFile(suffix='.onnx').name
ts_file = tempfile.NamedTemporaryFile(suffix='.pt').name
test_img = torch.rand(1, 3, 8, 8)
output_names = ['output']
input_names = ['input']
@pytest.mark.skip(reason='This a not test class but a utility class.')
class TestModel(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return x + test_img
model = TestModel().eval()
@pytest.fixture(autouse=True, scope='module')
def generate_onnx_file():
with torch.no_grad():
torch.onnx.export(
model,
test_img,
onnx_file,
output_names=output_names,
input_names=input_names,
keep_initializers_as_inputs=True,
do_constant_folding=True,
verbose=False,
opset_version=11,
dynamic_axes=None)
@pytest.fixture(autouse=True, scope='module')
def generate_torchscript_file():
import mmcv
backend = Backend.TORCHSCRIPT.value
deploy_cfg = mmcv.Config({'backend_config': dict(type=backend)})
from mmdeploy.apis.torch_jit import trace
context_info = dict(deploy_cfg=deploy_cfg)
output_prefix = osp.splitext(ts_file)[0]
example_inputs = torch.rand(1, 3, 8, 8)
trace(
model,
example_inputs,
output_path_prefix=output_prefix,
backend=backend,
context_info=context_info)
def onnx2backend(backend, onnx_file):
if backend == Backend.TENSORRT:
from mmdeploy.backend.tensorrt import from_onnx
backend_file = tempfile.NamedTemporaryFile(suffix='.engine').name
from_onnx(
onnx_file,
osp.splitext(backend_file)[0], {
'input': {
'min_shape': [1, 3, 8, 8],
'opt_shape': [1, 3, 8, 8],
'max_shape': [1, 3, 8, 8]
}
})
return backend_file
elif backend == Backend.ONNXRUNTIME:
return onnx_file
elif backend == Backend.PPLNN:
from mmdeploy.apis.pplnn import onnx2pplnn
algo_file = tempfile.NamedTemporaryFile(suffix='.json').name
onnx2pplnn(algo_file=algo_file, onnx_model=onnx_file)
return onnx_file, algo_file
elif backend == Backend.NCNN:
from mmdeploy.backend.ncnn.init_plugins import get_onnx2ncnn_path
onnx2ncnn_path = get_onnx2ncnn_path()
param_file = tempfile.NamedTemporaryFile(suffix='.param').name
bin_file = tempfile.NamedTemporaryFile(suffix='.bin').name
subprocess.call([onnx2ncnn_path, onnx_file, param_file, bin_file])
return param_file, bin_file
elif backend == Backend.OPENVINO:
from mmdeploy.apis.openvino import from_onnx, get_output_model_file
backend_dir = tempfile.TemporaryDirectory().name
backend_file = get_output_model_file(onnx_file, backend_dir)
input_info = {'input': test_img.shape}
output_names = ['output']
work_dir = backend_dir
from_onnx(onnx_file, work_dir, input_info, output_names)
return backend_file
def create_wrapper(backend, model_files):
if backend == Backend.TENSORRT:
from mmdeploy.backend.tensorrt import TRTWrapper
trt_model = TRTWrapper(model_files, output_names)
return trt_model
elif backend == Backend.ONNXRUNTIME:
from mmdeploy.backend.onnxruntime import ORTWrapper
ort_model = ORTWrapper(model_files, 'cpu', output_names)
return ort_model
elif backend == Backend.PPLNN:
from mmdeploy.backend.pplnn import PPLNNWrapper
onnx_file, algo_file = model_files
pplnn_model = PPLNNWrapper(onnx_file, algo_file, 'cpu', output_names)
return pplnn_model
elif backend == Backend.NCNN:
from mmdeploy.backend.ncnn import NCNNWrapper
param_file, bin_file = model_files
ncnn_model = NCNNWrapper(param_file, bin_file, output_names)
return ncnn_model
elif backend == Backend.OPENVINO:
from mmdeploy.backend.openvino import OpenVINOWrapper
openvino_model = OpenVINOWrapper(model_files, output_names)
return openvino_model
elif backend == Backend.TORCHSCRIPT:
from mmdeploy.backend.torchscript import TorchscriptWrapper
torchscript_model = TorchscriptWrapper(
model_files, input_names=input_names, output_names=output_names)
return torchscript_model
else:
raise NotImplementedError(f'Unknown backend type: {backend.value}')
def run_wrapper(backend, wrapper, input):
if backend == Backend.TENSORRT:
input = input.cuda()
results = wrapper({'input': input})['output']
results = results.detach().cpu()
return results
elif backend == Backend.ONNXRUNTIME:
results = wrapper({'input': input})['output']
results = results.detach().cpu()
return results
elif backend == Backend.PPLNN:
results = wrapper({'input': input})['output']
results = results.detach().cpu()
return results
elif backend == Backend.NCNN:
input = input.float()
results = wrapper({'input': input})['output']
results = results.detach().cpu()
return results
elif backend == Backend.OPENVINO:
results = wrapper({'input': input})['output']
results = results.detach().cpu()
return results
elif backend == Backend.TORCHSCRIPT:
results = wrapper({'input': input})['output']
return results
else:
raise NotImplementedError(f'Unknown backend type: {backend.value}')
ALL_BACKEND = [
Backend.TENSORRT, Backend.ONNXRUNTIME, Backend.PPLNN, Backend.NCNN,
Backend.OPENVINO, Backend.TORCHSCRIPT
]
@pytest.mark.parametrize('backend', ALL_BACKEND)
def test_wrapper(backend):
check_backend(backend, True)
if backend == Backend.TORCHSCRIPT:
model_files = ts_file
else:
model_files = onnx2backend(backend, onnx_file)
assert model_files is not None
wrapper = create_wrapper(backend, model_files)
assert wrapper is not None
results = run_wrapper(backend, wrapper, test_img)
assert results is not None