220 lines
7.4 KiB
Python
220 lines
7.4 KiB
Python
import importlib
|
|
import os
|
|
import tempfile
|
|
|
|
import mmcv
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
import mmdeploy.apis.onnxruntime as ort_apis
|
|
import mmdeploy.apis.ppl as ppl_apis
|
|
import mmdeploy.apis.tensorrt as trt_apis
|
|
import mmdeploy.apis.test as api_test
|
|
import mmdeploy.apis.utils as api_utils
|
|
from mmdeploy.utils.constants import Backend, Codebase
|
|
from mmdeploy.utils.test import SwitchBackendWrapper
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
|
|
@pytest.mark.skipif(
|
|
not importlib.util.find_spec('tensorrt'), reason='requires tensorrt')
|
|
def test_TensorRTRestorer():
|
|
# force add backend wrapper regardless of plugins
|
|
from mmdeploy.apis.tensorrt.tensorrt_utils import TRTWrapper
|
|
trt_apis.__dict__.update({'TRTWrapper': TRTWrapper})
|
|
|
|
# simplify backend inference
|
|
outputs = {
|
|
'output': torch.rand(1, 3, 64, 64).cuda(),
|
|
}
|
|
|
|
with SwitchBackendWrapper(TRTWrapper) as wrapper:
|
|
wrapper.set(outputs=outputs)
|
|
|
|
from mmdeploy.mmedit.apis.inference import TensorRTRestorer
|
|
trt_restorer = TensorRTRestorer('', 0)
|
|
imgs = torch.rand(1, 3, 64, 64).cuda()
|
|
|
|
results = trt_restorer.forward(imgs)
|
|
assert results is not None, ('failed to get output using '
|
|
'TensorRTRestorer')
|
|
|
|
results = trt_restorer.forward(imgs, test_mode=True)
|
|
assert results is not None, ('failed to get output using '
|
|
'TensorRTRestorer')
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime')
|
|
def test_ONNXRuntimeRestorer():
|
|
# force add backend wrapper regardless of plugins
|
|
from mmdeploy.apis.onnxruntime.onnxruntime_utils import ORTWrapper
|
|
ort_apis.__dict__.update({'ORTWrapper': ORTWrapper})
|
|
|
|
# simplify backend inference
|
|
outputs = torch.rand(1, 3, 64, 64)
|
|
|
|
with SwitchBackendWrapper(ORTWrapper) as wrapper:
|
|
wrapper.set(outputs=outputs)
|
|
|
|
from mmdeploy.mmedit.apis.inference import ONNXRuntimeRestorer
|
|
ort_restorer = ONNXRuntimeRestorer('', 0)
|
|
imgs = torch.rand(1, 3, 64, 64)
|
|
|
|
results = ort_restorer.forward(imgs)
|
|
assert results is not None, 'failed to get output using '\
|
|
'ONNXRuntimeRestorer'
|
|
|
|
results = ort_restorer.forward(imgs, test_mode=True)
|
|
assert results is not None, 'failed to get output using '\
|
|
'ONNXRuntimeRestorer'
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
|
|
@pytest.mark.skipif(
|
|
not importlib.util.find_spec('pyppl'), reason='requires pyppl')
|
|
def test_PPLRestorer():
|
|
# force add backend wrapper regardless of plugins
|
|
from mmdeploy.apis.ppl.ppl_utils import PPLWrapper
|
|
ppl_apis.__dict__.update({'PPLWrapper': PPLWrapper})
|
|
|
|
# simplify backend inference
|
|
outputs = torch.rand(1, 3, 64, 64)
|
|
|
|
with SwitchBackendWrapper(PPLWrapper) as wrapper:
|
|
wrapper.set(outputs=outputs)
|
|
|
|
from mmdeploy.mmedit.apis.inference import PPLRestorer
|
|
ppl_restorer = PPLRestorer('', 0)
|
|
imgs = torch.rand(1, 3, 64, 64)
|
|
|
|
results = ppl_restorer.forward(imgs)
|
|
assert results is not None, 'failed to get output using PPLRestorer'
|
|
|
|
results = ppl_restorer.forward(imgs, test_mode=True)
|
|
assert results is not None, 'failed to get output using PPLRestorer'
|
|
|
|
|
|
model_cfg = 'tests/test_mmedit/data/model.py'
|
|
deploy_cfg = mmcv.Config(
|
|
dict(
|
|
backend_config=dict(type='onnxruntime'),
|
|
codebase_config=dict(type='mmedit', task='SuperResolution'),
|
|
onnx_config=dict(
|
|
type='onnx',
|
|
export_params=True,
|
|
keep_initializers_as_inputs=False,
|
|
opset_version=11,
|
|
input_shape=None,
|
|
input_names=['input'],
|
|
output_names=['output'])))
|
|
input_img = torch.rand(1, 3, 64, 64)
|
|
input = {'lq': input_img}
|
|
|
|
|
|
def test_init_pytorch_model():
|
|
model = api_utils.init_pytorch_model(
|
|
Codebase.MMEDIT, model_cfg=model_cfg, device='cpu')
|
|
assert model is not None
|
|
|
|
|
|
def create_backend_model():
|
|
if not importlib.util.find_spec('onnxruntime'):
|
|
pytest.skip('requires onnxruntime')
|
|
from mmdeploy.apis.onnxruntime.onnxruntime_utils import ORTWrapper
|
|
ort_apis.__dict__.update({'ORTWrapper': ORTWrapper})
|
|
|
|
# simplify backend inference
|
|
|
|
wrapper = SwitchBackendWrapper(ORTWrapper)
|
|
wrapper.set(model_cfg=model_cfg, deploy_cfg=deploy_cfg)
|
|
model = api_utils.init_backend_model([''], model_cfg, deploy_cfg)
|
|
|
|
return model, wrapper
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime')
|
|
def test_init_backend_model():
|
|
model, wrapper = create_backend_model()
|
|
assert model is not None
|
|
|
|
# Recovery
|
|
wrapper.recover()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime')
|
|
def test_run_inference():
|
|
model, wrapper = create_backend_model()
|
|
result = api_utils.run_inference(Codebase.MMEDIT, input, model)
|
|
assert isinstance(result, np.ndarray)
|
|
|
|
# Recovery
|
|
wrapper.recover()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime')
|
|
def test_visualize():
|
|
model, wrapper = create_backend_model()
|
|
result = api_utils.run_inference(Codebase.MMEDIT, input, model)
|
|
with tempfile.TemporaryDirectory() as dir:
|
|
filename = dir + 'tmp.jpg'
|
|
api_utils.visualize(Codebase.MMEDIT, input, result, model, filename,
|
|
Backend.ONNXRUNTIME)
|
|
assert os.path.exists(filename)
|
|
|
|
# Recovery
|
|
wrapper.recover()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not importlib.util.find_spec('onnxruntime'), reason='requires onnxruntime')
|
|
def test_inference_model():
|
|
numpy_img = np.random.rand(64, 64, 3)
|
|
with tempfile.TemporaryDirectory() as dir:
|
|
filename = dir + 'tmp.jpg'
|
|
model, wrapper = create_backend_model()
|
|
from mmdeploy.apis.inference import inference_model
|
|
inference_model(model_cfg, deploy_cfg, model, numpy_img, 'cpu',
|
|
Backend.ONNXRUNTIME, filename, False)
|
|
assert os.path.exists(filename)
|
|
|
|
# Recovery
|
|
wrapper.recover()
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
|
|
def test_test():
|
|
from mmcv.parallel import MMDataParallel
|
|
with tempfile.TemporaryDirectory() as dir:
|
|
|
|
# Export a complete model
|
|
numpy_img = np.random.rand(50, 50, 3)
|
|
onnx_filename = 'end2end.onnx'
|
|
onnx_path = os.path.join(dir, onnx_filename)
|
|
from mmdeploy.apis import torch2onnx
|
|
torch2onnx(numpy_img, dir, onnx_filename, deploy_cfg, model_cfg)
|
|
assert os.path.exists(onnx_path)
|
|
|
|
# Prepare dataloader
|
|
dataset = api_utils.build_dataset(
|
|
Codebase.MMEDIT, model_cfg, dataset_type='test')
|
|
assert dataset is not None, 'Failed to build dataset'
|
|
dataloader = api_utils.build_dataloader(Codebase.MMEDIT, dataset, 1, 1)
|
|
assert dataloader is not None, 'Failed to build dataloader'
|
|
|
|
# Prepare model
|
|
model = api_utils.init_backend_model([onnx_path], model_cfg,
|
|
deploy_cfg)
|
|
model = MMDataParallel(model, device_ids=[0])
|
|
assert model is not None
|
|
|
|
# Run test
|
|
outputs = api_test.single_gpu_test(Codebase.MMEDIT, model, dataloader)
|
|
assert outputs is not None
|
|
api_test.post_process_outputs(outputs, dataset, model_cfg,
|
|
Codebase.MMEDIT)
|