mmdeploy/tests/test_apis/test_onnx2openvino.py

# Copyright (c) OpenMMLab. All rights reserved.
import os
import os.path as osp
import tempfile

import mmcv
import numpy as np
import pytest
import torch
import torch.nn as nn

from mmdeploy.utils import Backend
from mmdeploy.utils.test import backend_checker, get_random_name


@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 * 0.5


def generate_onnx_file(model, export_img, onnx_file, input_name, output_name):
    with torch.no_grad():
        dynamic_axes = {
            input_name: {
                0: 'batch',
                2: 'width',
                3: 'height'
            },
            output_name: {
                0: 'batch'
            }
        }
        torch.onnx.export(
            model,
            export_img,
            onnx_file,
            output_names=[output_name],
            input_names=[input_name],
            keep_initializers_as_inputs=True,
            do_constant_folding=True,
            verbose=False,
            opset_version=11,
            dynamic_axes=dynamic_axes)
        assert osp.exists(onnx_file)


def get_outputs(pytorch_model, openvino_model_path, input, input_name,
                output_name):
    output_pytorch = pytorch_model(input).numpy()

    from mmdeploy.backend.openvino import OpenVINOWrapper
    openvino_model = OpenVINOWrapper(openvino_model_path)
    openvino_output = openvino_model({input_name: input})[output_name]

    return output_pytorch, openvino_output


@backend_checker(Backend.OPENVINO)
def test_onnx2openvino():
    from mmdeploy.apis.openvino import get_output_model_file, onnx2openvino
    pytorch_model = TestModel().eval()
    export_img = torch.rand([1, 3, 8, 8])
    onnx_file = tempfile.NamedTemporaryFile(suffix='.onnx').name
    input_name = get_random_name()
    output_name = get_random_name()
    generate_onnx_file(pytorch_model, export_img, onnx_file, input_name,
                       output_name)

    input_info = {input_name: export_img.shape}
    output_names = [output_name]
    openvino_dir = tempfile.TemporaryDirectory().name
    onnx2openvino(input_info, output_names, onnx_file, openvino_dir)
    openvino_model_path = get_output_model_file(onnx_file, openvino_dir)
    assert osp.exists(openvino_model_path), \
        'The file (.xml) for OpenVINO IR has not been created.'

    test_img = torch.rand([1, 3, 16, 16])
    output_pytorch, openvino_output = get_outputs(pytorch_model,
                                                  openvino_model_path,
                                                  test_img, input_name,
                                                  output_name)
    assert np.allclose(output_pytorch, openvino_output), \
        'OpenVINO and PyTorch outputs are not the same.'


@backend_checker(Backend.OPENVINO)
def test_can_not_run_onnx2openvino_without_mo():
    current_environ = dict(os.environ)
    os.environ.clear()

    is_error = False
    try:
        from mmdeploy.apis.openvino import onnx2openvino
        onnx2openvino({}, ['output'], 'tmp.onnx', '/tmp')
    except RuntimeError:
        is_error = True

    os.environ.update(current_environ)
    assert is_error, \
        'The onnx2openvino script was launched without checking for MO.'


@backend_checker(Backend.OPENVINO)
def test_get_input_info_from_cfg():
    from mmdeploy.apis.openvino import get_input_info_from_cfg

    # Test 1
    deploy_cfg = mmcv.Config()
    with pytest.raises(KeyError):
        get_input_info_from_cfg(deploy_cfg)

    # Test 2
    input_name = 'input'
    height, width = 600, 1000
    shape = [1, 3, height, width]
    expected_input_info = {input_name: shape}
    deploy_cfg = mmcv.Config({
        'backend_config': {
            'model_inputs': [{
                'opt_shapes': expected_input_info
            }]
        }
    })
    input_info = get_input_info_from_cfg(deploy_cfg)
    assert input_info == expected_input_info, 'Test 2: ' \
        'The expected value of \'input_info\' does not match the received one.'

    # Test 3
    # The case where the input name in 'onnx_config'
    # is different from 'backend_config'.
    onnx_config_input_name = get_random_name(1234)
    deploy_cfg.merge_from_dict(
        {'onnx_config': {
            'input_names': [onnx_config_input_name]
        }})
    expected_input_info = {onnx_config_input_name: shape}
    input_info = get_input_info_from_cfg(deploy_cfg)
    assert input_info == expected_input_info, 'Test 3: ' \
        'The expected value of \'input_info\' does not match the received one.'

    # Test 4
    # The case where 'backend_config.model_inputs.opt_shapes'
    # is given by a list, not a dictionary.
    deploy_cfg.merge_from_dict(
        {'backend_config': {
            'model_inputs': [{
                'opt_shapes': [shape]
            }]
        }})
    input_info = get_input_info_from_cfg(deploy_cfg)
    assert input_info == expected_input_info, 'Test 4: ' \
        'The expected value of \'input_info\' does not match the received one.'