Merge pull request fastmachinelearning#67 from i-colbert/feat/quant_s…

…ubpixel_to_deconv

Feat: Adding quantization support for sub-pixel convolution to deconvolution transformation

src/qonnx/transformation/subpixel_to_deconv.py

@@ -111,8 +111,30 @@ def apply(self, model):
                     group = get_by_name(n.attribute, "group").i
                     if group != 1:
                         warnings.warn("Skipping sub-pixel conv with group > 1. Not yet supported.")
+
+                    # The weights of the convolution can be generated by another input op if the model is
+                    # quantized. Preliminary support for quantization focuses on QONNX ops (i.e., Quant)
                     weight_name = n.input[1]
-                    W_conv = model.get_initializer(weight_name)  # (OC, IC, KH, KW)
+                    weight_prod = model.find_producer(weight_name)
+
+                    # If the producer is None, then it is initialized by the Conv node
+                    if weight_prod is None:
+                        W_conv = model.get_initializer(weight_name)  # (OC, IC, KH, KW)
+
+                    # If the convolution weights are not initialized by the convolution, then we need to
+                    # find the node is producing the weights
+                    else:
+                        if weight_prod.op_type == "Quant":
+                            [q_w_name, q_s_name, _, _] = weight_prod.input
+                            W_conv = model.get_initializer(q_w_name)
+                            W_scale = model.get_initializer(q_s_name)
+                            assert W_scale.ndim == 0, "Only supporting per-tensor quantization with this transformation."
+                        else:
+                            warnings.warn(
+                                f"Weight producer is {weight_prod.op_type}, not a QONNX Quant node. Not yet supported."
+                            )
+                            continue
+
                     kshape = get_by_name(n.attribute, "kernel_shape").ints
                     ifm_ch = model.get_tensor_shape(n.input[0])[1]  # assume NCHW
                     ofm_ch = model.get_tensor_shape(n.output[0])[1]  # assume NCHW
@@ -184,7 +206,11 @@ def apply(self, model):
                         strides=[block_size, block_size],
                         pads=deconv_pad,
                     )
-                    model.set_initializer(weight_name, W_deconv)
+                    W_deconv_init = weight_name
+                    if weight_prod is not None:
+                        W_deconv_init = q_w_name
+                    model.set_initializer(W_deconv_init, W_deconv)
+                    model.set_tensor_shape(weight_name, list(W_deconv.shape))
                     graph.node.insert(node_ind, deconv_node)
                     # remove old nodes
                     graph.node.remove(n)

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
@@ -44,8 +46,8 @@
 np.random.seed(0)
 
 
-def test_subpixel_to_deconv_espcn():
-    raw_m = get_data("qonnx.data", "onnx/bsd300x3-espcn/model.onnx")
+def test_subpixel_to_deconv_float_espcn():
+    raw_m = get_data("qonnx.data", "onnx/bsd300x3-espcn/float_model.onnx")
     model = ModelWrapper(raw_m)
     model = model.transform(InferShapes())
     iname = model.graph.input[0].name
@@ -57,9 +59,39 @@ def test_subpixel_to_deconv_espcn():
     new_model = model.transform(SubPixelToDeconvolution())
     # 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
+    assert "DepthToSpace" not in op_types, "Error: the DepthToSpace nodes would be removed."
     produced = oxe.execute_onnx(new_model, input_dict)[oname]
-    assert np.isclose(expected, produced, atol=1e-4).all()
+    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(