diff --git a/tests/transformation/test_subpixel_to_deconv.py b/tests/transformation/test_subpixel_to_deconv.py index 821bc398..70db533c 100644 --- a/tests/transformation/test_subpixel_to_deconv.py +++ b/tests/transformation/test_subpixel_to_deconv.py @@ -29,7 +29,9 @@ import pytest import numpy as np +import onnx import onnx.helper as oh +import onnx.numpy_helper as nph from onnx import TensorProto from onnx.checker import check_model from pkgutil import get_data @@ -62,6 +64,36 @@ def test_subpixel_to_deconv_float_espcn(): assert np.isclose(expected, produced, atol=1e-4).all(), "Error: expected output does not match the produced output." +def test_subpixel_to_deconv_quant_espcn(): + # get raw quantized model with reference input + raw_i = get_data("qonnx.data", "onnx/bsd300x3-espcn/test_data/input_0.pb") + raw_m = get_data("qonnx.data", "onnx/bsd300x3-espcn/quant_model.onnx") + # create model from the onnx file and infer the shapes + model = ModelWrapper(raw_m) + model = model.transform(InferShapes()) + iname = model.graph.input[0].name + oname = model.graph.output[0].name + ishape = model.get_tensor_shape(iname) + # load the reference input tensor + input_tensor = onnx.load_tensor_from_string(raw_i) + input_tensor = nph.to_array(input_tensor) + assert list(input_tensor.shape) == ishape, "Error: reference input doesn't match loaded model." + input_dict = {iname: input_tensor} + # get the output from the sub-pixel convolution model + output_subpixel_conv = oxe.execute_onnx(model, input_dict)[oname] + # translate the sub-pixel convolution to the deconvolution + new_model = model.transform(SubPixelToDeconvolution()) + new_model = new_model.transform(InferShapes()) + # check that there are no DepthToSpace ops left + op_types = list(map(lambda x: x.op_type, new_model.graph.node)) + assert "DepthToSpace" not in op_types, "Error: the DepthToSpace nodes would be removed." + # get the output from the deconvolution model + output_deconv = oxe.execute_onnx(new_model, input_dict)[oname] + assert np.isclose( + output_deconv, output_subpixel_conv, atol=1 / 255.0, rtol=1 / 255.0 + ).all(), "Error: expected output does not match the produced output." + + def create_subpixel_conv_model( in_channels: int, out_channels: int, input_dim: int, kernel_size: int, upscale_factor: int, bias: bool = False ):