diff --git a/tests/test_ops/__init__.py b/tests/test_ops/__init__.py index dcf7ee9ae..4660224a2 100644 --- a/tests/test_ops/__init__.py +++ b/tests/test_ops/__init__.py @@ -1,3 +1,3 @@ -from .utils import TestOnnxRTExporter, TestTensorRTExporter +from .utils import TestNCNNExporter, TestOnnxRTExporter, TestTensorRTExporter -__all__ = ['TestTensorRTExporter', 'TestOnnxRTExporter'] +__all__ = ['TestTensorRTExporter', 'TestOnnxRTExporter', 'TestNCNNExporter'] diff --git a/tests/test_ops/test_ops.py b/tests/test_ops/test_ops.py index 92197238a..caf0ebc26 100644 --- a/tests/test_ops/test_ops.py +++ b/tests/test_ops/test_ops.py @@ -4,10 +4,11 @@ import torch.nn as nn from mmdeploy.core import register_extra_symbolics from mmdeploy.utils.test import WrapFunction -from .utils import TestOnnxRTExporter, TestTensorRTExporter +from .utils import TestNCNNExporter, TestOnnxRTExporter, TestTensorRTExporter TEST_ONNXRT = TestOnnxRTExporter() TEST_TENSORRT = TestTensorRTExporter() +TEST_NCNN = TestNCNNExporter() @pytest.mark.parametrize('backend', [TEST_TENSORRT, TEST_ONNXRT]) @@ -185,3 +186,49 @@ def test_instance_norm(backend, dynamic_axes=dynamic_axes, output_names=['output'], save_dir=save_dir) + + +@pytest.mark.parametrize('backend', [TEST_NCNN]) +@pytest.mark.parametrize('k', [1, 3, 5]) +@pytest.mark.parametrize('dim', [1, 2, 3]) +@pytest.mark.parametrize('largest', [True, False]) +@pytest.mark.parametrize('sorted', [True, False]) +def test_topk(backend, + k, + dim, + largest, + sorted, + input_list=None, + save_dir=None): + backend.check_env() + + if not input_list: + input = torch.rand(1, 8, 12, 17) + assert input.shape[0] == 1, (f'ncnn batch must be 1, \ + but not {input.shape[0]}') + cfg = dict() + register_extra_symbolics(cfg=cfg, opset=11) + + def wrapped_function(inputs): + return torch.Tensor.topk(inputs, k, dim, largest, sorted) + + wrapped_model = WrapFunction(wrapped_function) + + # when the 'sorted' attribute is False, pytorch will return + # a hard to expect result, which only features that the topk + # number is right. So the Topk unittest only check whether the + # topk elements are right, all the possible order will be accepted. + if not sorted: + backend.run_and_validate( + wrapped_model, [input.float()], + 'topk' + f'_no_sorted_dim_{dim}', + input_names=['inputs'], + output_names=['data', 'index'], + save_dir=save_dir) + else: + backend.run_and_validate( + wrapped_model, [input.float()], + 'topk', + input_names=['inputs'], + output_names=['data', 'index'], + save_dir=save_dir) diff --git a/tests/test_ops/utils.py b/tests/test_ops/utils.py index 2ee43444e..cf73ac185 100644 --- a/tests/test_ops/utils.py +++ b/tests/test_ops/utils.py @@ -1,4 +1,5 @@ import os +import subprocess import tempfile import mmcv @@ -6,6 +7,7 @@ import onnx import pytest import torch +import mmdeploy.apis.ncnn as ncnn_apis import mmdeploy.apis.onnxruntime as ort_apis import mmdeploy.apis.tensorrt as trt_apis from mmdeploy.utils.test import assert_allclose @@ -149,3 +151,77 @@ class TestTensorRTExporter: trt_outputs[name].cpu().float() for name in output_names ] assert_allclose(model_outputs, trt_outputs, tolerate_small_mismatch) + + +@pytest.mark.skip(reason='This a not test class but a utility class.') +class TestNCNNExporter: + + def __init__(self): + self.backend_name = 'ncnn' + + def check_env(self): + if not ncnn_apis.is_available(): + pytest.skip( + 'NCNN is not installed or custom ops are not compiled.') + + def run_and_validate(self, + model, + inputs_list, + model_name='tmp', + tolerate_small_mismatch=False, + do_constant_folding=True, + dynamic_axes=None, + output_names=None, + input_names=None, + save_dir=None): + if not save_dir: + onnx_file_path = tempfile.NamedTemporaryFile().name + ncnn_param_path = tempfile.NamedTemporaryFile().name + ncnn_bin_path = tempfile.NamedTemporaryFile().name + else: + onnx_file_path = os.path.join(save_dir, model_name + '.onnx') + ncnn_param_path = os.path.join(save_dir, model_name + '.param') + ncnn_bin_path = os.path.join(save_dir, model_name + '.bin') + + with torch.no_grad(): + torch.onnx.export( + model, + tuple(inputs_list), + onnx_file_path, + export_params=True, + keep_initializers_as_inputs=True, + input_names=input_names, + output_names=output_names, + do_constant_folding=do_constant_folding, + dynamic_axes=dynamic_axes, + opset_version=11) + + onnx2ncnn_path = ncnn_apis.get_onnx2ncnn_path() + subprocess.call( + [onnx2ncnn_path, onnx_file_path, ncnn_param_path, ncnn_bin_path]) + + with torch.no_grad(): + model_outputs = model(*inputs_list) + if isinstance(model_outputs, torch.Tensor): + model_outputs = [model_outputs] + else: + model_outputs = list(model_outputs) + model_outputs = [ + model_output.float() for model_output in model_outputs + ] + ncnn_model = ncnn_apis.NCNNWrapper(ncnn_param_path, ncnn_bin_path, + output_names) + ncnn_outputs = ncnn_model(dict(zip(input_names, inputs_list))) + ncnn_outputs = [ncnn_outputs[name] for name in output_names] + + if model_name.startswith('topk_no_sorted'): + dim = int(model_name.split('_')[-1]) + model_outputs = torch.stack(model_outputs, dim=-1).\ + sort(dim=dim).values + ncnn_outputs = torch.stack(ncnn_outputs, dim=-1).\ + sort(dim=dim).values + assert_allclose([model_outputs], [ncnn_outputs], + tolerate_small_mismatch) + else: + assert_allclose(model_outputs, ncnn_outputs, + tolerate_small_mismatch)