[LowerConv] support lowering Conv with Quant node on weights

Author: maltanar

src/qonnx/transformation/lower_convs_to_matmul.py

@@ -51,10 +51,6 @@ def apply(self, model):
                 warnings.warn("Found Conv node with bias, skipping")
                 continue
 
-            if model.get_initializer(node.input[1]) is None:
-                warnings.warn("Found Conv node with non-initialized weight, skipping")
-                continue
-
             # extract parameters of node
             (
                 cnv_input,
@@ -67,6 +63,7 @@ def apply(self, model):
                 stride_w,
                 group,
                 weight_name,
+                conv_weight_inp_name,
                 W_conv,
                 ifm_ch,
                 ofm_ch,
@@ -78,6 +75,10 @@ def apply(self, model):
                 pad,
             ) = self.extract_conv_params(model, node)
 
+            if W_conv is None:
+                warnings.warn("Found Conv node with non-initialized weight, skipping")
+                continue
+
             # if depthwise conv create sparse matrix and variable "dw"
             # to store as attribute in Im2Col that indicates that the created
             # Im2Col node belongs to a depthwise convolution
@@ -108,6 +109,8 @@ def apply(self, model):
             # transpose to get ONNX-compatible [k_h*k_w*IFM][OFM] matrix
             W_matmul = W_matmul.T
             model.set_initializer(weight_name, W_matmul)
+            if weight_name != conv_weight_inp_name:
+                model.set_tensor_shape(conv_weight_inp_name, W_matmul.shape)
 
             # create new intermediate values
             inp_trans_out = helper.make_tensor_value_info(
@@ -158,7 +161,7 @@ def apply(self, model):
 
             matmul_input = im2col_out if need_im2col else inp_trans_out
             # do matmul
-            matmul_node = helper.make_node("MatMul", [matmul_input, weight_name], [matmul_out])
+            matmul_node = helper.make_node("MatMul", [matmul_input, conv_weight_inp_name], [matmul_out])
             # NHWC -> NCHW
             out_trans_node = helper.make_node("Transpose", [matmul_out], [cnv_output], perm=[0, 3, 1, 2])
 
@@ -182,7 +185,14 @@ def extract_conv_params(self, model, node):
         stride_w = get_by_name(node.attribute, "strides").ints[1]
         group = get_by_name(node.attribute, "group").i
         weight_name = node.input[1]
+        conv_weight_inp_name = node.input[1]
         W_conv = model.get_initializer(weight_name)
+        if W_conv is None:
+            # check to see if there is an immediate quantizer node feeding the weight input
+            w_producer = model.find_producer(weight_name)
+            if not (w_producer is None) and w_producer.op_type == "Quant":
+                W_conv = model.get_initializer(w_producer.input[0])
+                weight_name = w_producer.input[0]
         ifm_ch = model.get_tensor_shape(cnv_input)[1]  # assume NCHW
         ofm_ch = model.get_tensor_shape(cnv_output)[1]  # assume NCHW
         ifm_dim_h = model.get_tensor_shape(cnv_input)[2]  # assume NCHW
@@ -217,6 +227,7 @@ def extract_conv_params(self, model, node):
             stride_w,
             group,
             weight_name,
+            conv_weight_inp_name,
             W_conv,
             ifm_ch,
             ofm_ch,