refactor LowerConvsToMatMul class, increase rtol in test_conv_lowering_convmnist()

makoeppel · makoeppel · commit 1b2774c06354 · 2024-06-17T09:28:36.000+02:00
diff --git a/README.md b/README.md
@@ -125,6 +125,7 @@ source venv/bin/activate
 pip install -e .[qkeras,testing]
 ```
 
+### Test suite
 Run entire test suite, parallelized across CPU cores:
 ```
 pytest -n auto --verbose
diff --git a/docs/index.rst b/docs/index.rst
@@ -63,6 +63,9 @@ Install in editable mode in a venv:
    pip install -e .[testing, docs, notebooks]
 
 
+Test suite
+++++++++++
+
 Run entire test suite, parallelized across CPU cores:
 
 ::
diff --git a/src/qonnx/transformation/lower_convs_to_matmul.py b/src/qonnx/transformation/lower_convs_to_matmul.py
@@ -42,167 +42,143 @@ class LowerConvsToMatMul(Transformation):
     def apply(self, model):
         model = model.transform(ExtractBiasFromConv())
         graph = model.graph
-        node_ind = 0
         graph_modified = False
-        for n in graph.node:
-            node_ind += 1
-            if n.op_type == "Conv":
-                if len(n.input) == 3:
-                    warnings.warn("Found Conv node with bias, skipping")
-                    continue
-                cnv_input = n.input[0]
-                cnv_output = n.output[0]
-                idt = model.get_tensor_datatype(cnv_input)
-                odt = model.get_tensor_datatype(cnv_output)
-                # extract conv parameters
-                k = get_by_name(n.attribute, "kernel_shape").ints
-                k_h = k[0]
-                k_w = k[1]
-                stride_h = get_by_name(n.attribute, "strides").ints[0]
-                stride_w = get_by_name(n.attribute, "strides").ints[1]
-                group = get_by_name(n.attribute, "group").i
-                weight_name = n.input[1]
-                W_conv = model.get_initializer(weight_name)
-                ifm_ch = model.get_tensor_shape(n.input[0])[1]  # assume NCHW
-                ofm_ch = model.get_tensor_shape(n.output[0])[1]  # assume NCHW
-                ifm_dim_h = model.get_tensor_shape(n.input[0])[2]  # assume NCHW
-                ifm_dim_w = model.get_tensor_shape(n.input[0])[3]
-                ofm_dim_h = model.get_tensor_shape(n.output[0])[2]  # assume NCHW
-                ofm_dim_w = model.get_tensor_shape(n.output[0])[3]
-                dilation_attr = get_by_name(n.attribute, "dilations")
-                if dilation_attr is not None:
-                    dilation = dilation_attr.ints
-                else:
-                    dilation = [1, 1]  # default value
-                # handle both auto_pad and explicit padding
-                auto_pad = get_by_name(n.attribute, "auto_pad")
-                if auto_pad is not None:
-                    # find equivalent specified padding
-                    auto_pad = auto_pad.s.decode("utf-8")
-                    if auto_pad == "NOTSET":
-                        # use specified padding
-                        pad = get_by_name(n.attribute, "pads").ints
-                    else:
-                        pad = auto_pad_to_explicit_padding(
-                            auto_pad,
-                            ifm_dim_h,
-                            ifm_dim_w,
-                            k_h,
-                            k_w,
-                            stride_h,
-                            stride_w,
-                            len(model.get_tensor_shape(n.input[0])) - 2,
-                        )
-                else:
-                    # use specified padding
-                    pad = get_by_name(n.attribute, "pads").ints
-
-                # If len(pad) == 2, assume no padding for other dimension
-                if len(pad) == 2:  # only one dimension should be padded
-                    assert ifm_dim_h == 1 or ifm_dim_w == 1, "Padding is assumed to be 1D, image is 2D"
-
-                # if depthwise conv create sparse matrix and variable "dw"
-                # to store as attribute in Im2Col that indicates that the created
+        for node_ind, node in enumerate(graph.node, start=1):
+            if node.op_type != "Conv":
+                continue
+
+            if len(node.input) == 3:
+                warnings.warn("Found Conv node with bias, skipping")
+                continue
+
+            # extract parameters of node
+            (cnv_input, cnv_output, cnv_input_datatype, cnv_output_datatype,
+             k_h, k_w, stride_h, stride_w, group, weight_name, W_conv, ifm_ch,
+             ofm_ch, ifm_dim_h, ifm_dim_w, ofm_dim_h, ofm_dim_w, dilation, pad) =\
+                self.extract_conv_params(model, node)
+
+            # 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
+            dw = False
+            if group == ifm_ch and ofm_ch == ifm_ch:
+                W_sparse = np.zeros((ofm_ch, ifm_ch, k_h, k_w))  # (OFM, IFM, k_H, k_W)
+                for ch in range(ifm_ch):
+                    W_sparse[ch][ch] = W_conv[ch][0]  # W_conv = [OFM, IFM, k_H, k_W]
+                W_conv = W_sparse.astype(np.float32)
+                # we need to store information of the
+                # sparsity of the weight matrix. For this
+                # we use the sparsity annotation of the
+                # weight tensor
+                sparsity = {"dw": {"kernel_shape": [k_h, k_w]}}
+                model.set_tensor_sparsity(weight_name, sparsity)
+                # additionally create variable "dw" to store
+                # as attribute in Im2Col that indicates that the created
                 # Im2Col node belongs to a depthwise convolution
-                dw = False
-                if group == ifm_ch and ofm_ch == ifm_ch:
-                    W_sparse = np.zeros((ofm_ch, ifm_ch, k_h, k_w))  # (OFM, IFM, k_H, k_W)
-                    for ch in range(ifm_ch):
-                        W_sparse[ch][ch] = W_conv[ch][0]  # W_conv = [OFM, IFM, k_H, k_W]
-                    W_conv = W_sparse.astype(np.float32)
-                    # we need to store information of the
-                    # sparsity of the weight matrix. For this
-                    # we use the sparsity annotation of the
-                    # weight tensor
-                    sparsity = {"dw": {"kernel_shape": [k_h, k_w]}}
-                    model.set_tensor_sparsity(weight_name, sparsity)
-                    # additionally create variable "dw" to store
-                    # as attribute in Im2Col that indicates that the created
-                    # Im2Col node belongs to a depthwise convolution
-                    dw = True
-
-                # reuse conv weights for new matmul weights
-                # conv weights are [OFM][IFM][k][k]
-                # first convert to [OFM][k][k][IFM] (to remain compatible with
-                # finn-hlslib and how it does im2col/sliding window)
-                W_matmul = W_conv.transpose(0, 2, 3, 1)  # W_conv = [OFM, IFM, k_H, k_W]
-                # reshape into [OFM][k*k*IFM] matrix
-                W_matmul = W_matmul.reshape(ofm_ch, ifm_ch * k_h * k_w)
-                # transpose to get ONNX-compatible [k*k*IFM][OFM] matrix
-                W_matmul = W_matmul.T
-                model.set_initializer(weight_name, W_matmul)
-
-                # create new intermediate values
-                inp_trans_out = helper.make_tensor_value_info(
-                    model.make_new_valueinfo_name(),
-                    TensorProto.FLOAT,
-                    (1, ifm_dim_h, ifm_dim_w, ifm_ch),  # NHWC
+                dw = True
+
+            # reuse conv weights for new matmul weights
+            # conv weights are [OFM][IFM][k][k]
+            # first convert to [OFM][k_h][k_w][IFM] (to remain compatible with
+            # finn-hlslib and how it does im2col/sliding window)
+            W_matmul = W_conv.transpose(0, 2, 3, 1)  # W_conv = [OFM, IFM, k_H, k_W]
+            # reshape into [OFM][k_h*k_w*IFM] matrix
+            W_matmul = W_matmul.reshape(ofm_ch, ifm_ch * k_h * k_w)
+            # transpose to get ONNX-compatible [k_h*k_w*IFM][OFM] matrix
+            W_matmul = W_matmul.T
+            model.set_initializer(weight_name, W_matmul)
+
+            # create new intermediate values
+            inp_trans_out = helper.make_tensor_value_info(
+                model.make_new_valueinfo_name(),
+                TensorProto.FLOAT,
+                (1, ifm_dim_h, ifm_dim_w, ifm_ch),  # NHWC
+            )
+            graph.value_info.append(inp_trans_out)
+            inp_trans_out = inp_trans_out.name
+            model.set_tensor_datatype(inp_trans_out, cnv_input_datatype)
+
+            # k_h=k_w==1: pointwise convolution, thus no im2col needed
+            need_im2col = any(p != 0 for p in pad) or k_h != 1 or k_w != 1 or stride_h != 1 or stride_w != 1
+
+            # create new intermediate values
+            matmul_out = helper.make_tensor_value_info(
+                model.make_new_valueinfo_name(), TensorProto.FLOAT, (1, ofm_dim_h, ofm_dim_w, ofm_ch)
+            )
+            graph.value_info.append(matmul_out)
+            matmul_out = matmul_out.name
+            model.set_tensor_datatype(matmul_out, cnv_output_datatype)
+
+            # create new nodes
+            # NCHW -> NHWC
+            inp_trans_node = helper.make_node("Transpose", [cnv_input], [inp_trans_out], perm=[0, 2, 3, 1])
+            nodes_to_insert = [inp_trans_node]
+
+            if need_im2col:
+                im2col_out = helper.make_tensor_value_info(
+                    model.make_new_valueinfo_name(), TensorProto.FLOAT, (1, ofm_dim_h, ofm_dim_w, ifm_ch * k_h * k_w)
                 )
-                graph.value_info.append(inp_trans_out)
-                inp_trans_out = inp_trans_out.name
-                model.set_tensor_datatype(inp_trans_out, idt)
-
-                need_im2col = True
-                if all(p == 0 for p in pad):
-                    padding = 0
-
-                # k_h=k_w==1: pointwise convolution, thus no im2col needed
-                if k_h == 1 and k_w == 1 and padding == 0 and stride_h == 1 and stride_w == 1:
-                    need_im2col = False
-
-                if need_im2col:
-                    im2col_out = helper.make_tensor_value_info(
-                        model.make_new_valueinfo_name(),
-                        TensorProto.FLOAT,
-                        (1, ofm_dim_h, ofm_dim_w, ifm_ch * k_h * k_w),
-                    )
-                    graph.value_info.append(im2col_out)
-                    im2col_out = im2col_out.name
-                    model.set_tensor_datatype(im2col_out, idt)
-
-                matmul_out = helper.make_tensor_value_info(
-                    model.make_new_valueinfo_name(),
-                    TensorProto.FLOAT,
-                    (1, ofm_dim_h, ofm_dim_w, ofm_ch),
+                graph.value_info.append(im2col_out)
+                im2col_out = im2col_out.name
+                model.set_tensor_datatype(im2col_out, cnv_input_datatype)
+                im2col_node = helper.make_node(
+                    "Im2Col", [inp_trans_out], [im2col_out], domain="qonnx.custom_op.general",
+                    stride=[stride_h, stride_w], kernel_size=[k_h, k_w], pad_amount=pad,
+                    input_shape="(1,{},{},{})".format(ifm_dim_h, ifm_dim_w, ifm_ch), depthwise=dw, dilations=dilation
                 )
-                graph.value_info.append(matmul_out)
-                matmul_out = matmul_out.name
-                model.set_tensor_datatype(matmul_out, odt)
-
-                # create new nodes
-                # NCHW -> NHWC
-                inp_trans_node = helper.make_node("Transpose", [cnv_input], [inp_trans_out], perm=[0, 2, 3, 1])
-                # lower input tensor
-                matmul_input = inp_trans_out
-                if need_im2col:
-                    matmul_input = im2col_out
-                    im2col_node = helper.make_node(
-                        "Im2Col",
-                        [inp_trans_out],
-                        [im2col_out],
-                        domain="qonnx.custom_op.general",
-                        stride=[stride_h, stride_w],
-                        kernel_size=[k_h, k_w],
-                        pad_amount=pad,
-                        input_shape="(1,{},{},{})".format(ifm_dim_h, ifm_dim_w, ifm_ch),
-                        depthwise=dw,
-                        dilations=dilation,
-                    )
-
-                # do matmul
-                matmul_node = helper.make_node("MatMul", [matmul_input, weight_name], [matmul_out])
-                # NHWC -> NCHW
-                out_trans_node = helper.make_node("Transpose", [matmul_out], [cnv_output], perm=[0, 3, 1, 2])
-                # insert nodes where the conv is to preserve topological ordering
-                graph.node.insert(node_ind, inp_trans_node)
-                if need_im2col:
-                    graph.node.insert(node_ind + 1, im2col_node)
-                    graph.node.insert(node_ind + 2, matmul_node)
-                    graph.node.insert(node_ind + 3, out_trans_node)
-                else:
-                    graph.node.insert(node_ind + 1, matmul_node)
-                    graph.node.insert(node_ind + 2, out_trans_node)
-                # remove old nodes
-                graph.node.remove(n)
+                nodes_to_insert.append(im2col_node)
+
+            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])
+            # NHWC -> NCHW
+            out_trans_node = helper.make_node("Transpose", [matmul_out], [cnv_output], perm=[0, 3, 1, 2])
+
+            nodes_to_insert.extend([matmul_node, out_trans_node])
+
+            # insert nodes where the conv is to preserve topological ordering
+            for i, insert_node in enumerate(nodes_to_insert):
+                graph.node.insert(node_ind + i, insert_node)
+            graph.node.remove(node)
 
         return (model, graph_modified)
+
+    def extract_conv_params(self, model, node):
+
+        cnv_input = node.input[0]
+        cnv_output = node.output[0]
+        cnv_input_datatype = model.get_tensor_datatype(cnv_input)
+        cnv_output_datatype = model.get_tensor_datatype(cnv_output)
+        k_h = get_by_name(node.attribute, "kernel_shape").ints[0]
+        k_w = get_by_name(node.attribute, "kernel_shape").ints[1]
+        stride_h = get_by_name(node.attribute, "strides").ints[0]
+        stride_w = get_by_name(node.attribute, "strides").ints[1]
+        group = get_by_name(node.attribute, "group").i
+        weight_name = node.input[1]
+        W_conv = model.get_initializer(weight_name)
+        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
+        ifm_dim_w = model.get_tensor_shape(cnv_input)[3]  # assume NCHW
+        ofm_dim_h = model.get_tensor_shape(cnv_output)[2]  # assume NCHW
+        ofm_dim_w = model.get_tensor_shape(cnv_output)[3]  # assume NCHW
+        dilation_attr = get_by_name(node.attribute, "dilations")
+        dilation = dilation_attr.ints if dilation_attr is not None else [1, 1]  # default value
+        auto_pad = get_by_name(node.attribute, "auto_pad")
+        if auto_pad is not None:
+            auto_pad = auto_pad.s.decode("utf-8")
+            if auto_pad == "NOTSET":
+                pad = get_by_name(node.attribute, "pads").ints
+            else:
+                pad = auto_pad_to_explicit_padding(
+                    auto_pad, ifm_dim_h, ifm_dim_w, k_h, k_w, stride_h, stride_w, len(model.get_tensor_shape(cnv_input)) - 2
+                )
+        else:
+            pad = get_by_name(node.attribute, "pads").ints
+
+        if len(pad) == 2:  # only one dimension should be padded
+            assert ifm_dim_h == 1 or ifm_dim_w == 1, "Padding is assumed to be 1D, image is 2D"
+
+        return (cnv_input, cnv_output, cnv_input_datatype, cnv_output_datatype, k_h, k_w, stride_h,
+                stride_w, group, weight_name, W_conv, ifm_ch, ofm_ch, ifm_dim_h, ifm_dim_w, ofm_dim_h,
+                ofm_dim_w, dilation, pad)
diff --git a/tests/transformation/test_conv_lowering.py b/tests/transformation/test_conv_lowering.py
@@ -65,7 +65,7 @@ def test_conv_lowering_convmnist():
     model = model.transform(InferShapes())
     output_dict_p = oxe.execute_onnx(model, input_dict)
     produced = output_dict_p[output_name]
-    assert np.isclose(produced, expected).all()
+    assert np.isclose(produced, expected, rtol=1.e-4).all()
 
 
 def run_conv_lowering_test(idt, k_h, k_w, ifm_dim_h, ifm_dim_w, ifm_ch, stride, padding, dilations, dw, bias):
