Merge pull request #631 from hleblevec/feature/make_scale_resize_nhwc

MakeScaleResizeNHWC tranformation

src/finn/transformation/streamline/reorder.py

@@ -725,6 +725,77 @@ class MakeMaxPoolNHWC(Transformation):
         return (model, graph_modified)
 
 
+class MakeScaleResizeNHWC(Transformation):
+    """
+    Converts the inputs and outputs for all scales Resize and Upsample nodes
+    from NCHW to NHWC.
+    """
+
+    def apply(self, model):
+        graph = model.graph
+        node_ind = 0
+        for n in graph.node:
+            node_ind += 1
+            if n.op_type == "Upsample" or n.op_type == "Resize":
+                if model.get_tensor_layout(n.input[0]) != DataLayout.NCHW:
+                    warnings.warn(
+                        "%s: Input not NCHW. Can't operate transformation on node."
+                        % n.name
+                    )
+                    continue
+                consumer = model.find_consumer(n.output[0])
+                producer = model.find_producer(n.input[0])
+                if n.op_type == "Upsample":
+                    scales_ind = 1
+                else:
+                    scales_ind = 2
+                if producer is not None and producer.op_type == "Transpose":
+                    perms = list(get_by_name(producer.attribute, "perm").ints)
+                    if perms == [0, 3, 1, 2]:
+                        old_value = model.get_initializer(n.input[scales_ind])
+                        new_value = np.array(
+                            [old_value[idx] for idx in (0, 2, 3, 1)],
+                            dtype=np.dtype("float32"),
+                        )
+                        model.set_initializer(n.input[scales_ind], new_value)
+                        start_name = producer.input[0]
+                        mid_name = n.input[0]
+                        end_name = n.output[0]
+                        (b, hi, wi, c) = model.get_tensor_shape(start_name)
+                        (b, c, ho, wo) = model.get_tensor_shape(end_name)
+                        producer.input[0] = mid_name
+                        producer.output[0] = end_name
+                        n.input[0] = start_name
+                        n.output[0] = mid_name
+                        model.set_tensor_shape(mid_name, (b, ho, wo, c))
+                        model.set_tensor_shape(end_name, (b, c, ho, wo))
+                        graph.node.remove(producer)
+                        graph.node.insert(node_ind, producer)
+                elif consumer is not None and consumer.op_type == "Transpose":
+                    perms = list(get_by_name(consumer.attribute, "perm").ints)
+                    if perms == [0, 2, 3, 1]:
+                        old_value = model.get_initializer(n.input[scales_ind])
+                        new_value = np.array(
+                            [old_value[idx] for idx in (0, 2, 3, 1)],
+                            dtype=np.dtype("float32"),
+                        )
+                        model.set_initializer(n.input[scales_ind], new_value)
+                        start_name = n.input[0]
+                        mid_name = consumer.input[0]
+                        end_name = consumer.output[0]
+                        (b, c, hi, wi) = model.get_tensor_shape(start_name)
+                        (b, c, ho, wo) = model.get_tensor_shape(mid_name)
+                        consumer.input[0] = start_name
+                        consumer.output[0] = mid_name
+                        n.input[0] = mid_name
+                        n.output[0] = end_name
+                        model.set_tensor_shape(mid_name, (b, hi, wi, c))
+                        model.set_tensor_shape(end_name, (b, ho, wo, c))
+                        graph.node.remove(consumer)
+                        graph.node.insert(node_ind - 1, consumer)
+        return (model, False)
+
+
 class MoveOpPastFork(Transformation):
     """Move node operations past graph forks. Used when a node before a fork
     can be merged with nodes in the branches

tests/transformation/streamline/test_scale_resize_nhwc.py

+import pytest
+
+import numpy as np
+import onnx
+import onnx.helper as oh
+import qonnx.core.data_layout as DataLayout
+from onnx import TensorProto
+from qonnx.core.datatype import DataType
+from qonnx.core.modelwrapper import ModelWrapper
+from qonnx.transformation.infer_data_layouts import InferDataLayouts
+from qonnx.transformation.infer_shapes import InferShapes
+from qonnx.util.basic import gen_finn_dt_tensor
+
+import finn.core.onnx_exec as oxe
+from finn.transformation.streamline.reorder import MakeScaleResizeNHWC
+
+
+def create_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
+    ofm_dim_h = ifm_dim[0] * scales[2]
+    ofm_dim_w = ifm_dim[1] * scales[3]
+    inp = oh.make_tensor_value_info(
+        "inp", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]]
+    )
+
+    param = oh.make_tensor_value_info("scales", TensorProto.FLOAT, [4])
+
+    # Not actually used, only needed for compliance with the Resize node interface
+    roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])
+
+    outp_up = oh.make_tensor_value_info(
+        "outp_up", TensorProto.FLOAT, [1, ifm_ch, ofm_dim_h, ofm_dim_w]
+    )
+    outp = oh.make_tensor_value_info(
+        "outp", TensorProto.FLOAT, [1, ofm_dim_h, ofm_dim_w, ifm_ch]
+    )
+
+    resize_node = oh.make_node(
+        "Resize",
+        inputs=["inp", "roi", "scales"],
+        outputs=["outp_up"],
+        name="Resize1",
+        mode=mode,
+    )
+
+    transpose_node = onnx.helper.make_node(
+        "Transpose",
+        inputs=["outp_up"],
+        outputs=["outp"],
+        name="Transpose1",
+        perm=[0, 2, 3, 1],
+    )
+
+    graph = oh.make_graph(
+        nodes=[resize_node, transpose_node],
+        name="resize_graph",
+        inputs=[inp],
+        outputs=[outp],
+        value_info=[outp_up, param, roi],
+    )
+
+    model = oh.make_model(graph, producer_name="resize_model1")
+    model = ModelWrapper(model)
+    model.set_tensor_datatype("inp", idt)
+    model.set_tensor_datatype("outp", idt)
+
+    model.set_tensor_layout("inp", DataLayout.NCHW)
+    model = model.transform(InferShapes())
+    model = model.transform(InferDataLayouts())
+
+    return model
+
+
+def create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt):
+    ofm_dim_h = ifm_dim[0] * scales[2]
+    ofm_dim_w = ifm_dim[1] * scales[3]
+    inp = oh.make_tensor_value_info(
+        "inp", TensorProto.FLOAT, [1, ifm_dim[0], ifm_dim[1], ifm_ch]
+    )
+
+    param = oh.make_tensor_value_info("scales", TensorProto.FLOAT, [4])
+
+    # Not actually used, only needed for compliance with the Resize node interface
+    roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])
+
+    outp = oh.make_tensor_value_info(
+        "outp", TensorProto.FLOAT, [1, ifm_ch, ofm_dim_h, ofm_dim_w]
+    )
+    outp_tr = oh.make_tensor_value_info(
+        "outp_tr", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]]
+    )
+
+    transpose_node = onnx.helper.make_node(
+        "Transpose",
+        inputs=["inp"],
+        outputs=["outp_tr"],
+        name="Transpose1",
+        perm=[0, 3, 1, 2],
+    )
+
+    resize_node = oh.make_node(
+        "Resize",
+        inputs=["outp_tr", "roi", "scales"],
+        outputs=["outp"],
+        name="Resize1",
+        mode=mode,
+    )
+
+    graph = oh.make_graph(
+        nodes=[transpose_node, resize_node],
+        name="resize_graph",
+        inputs=[inp],
+        outputs=[outp],
+        value_info=[outp_tr, param, roi],
+    )
+
+    model = oh.make_model(graph, producer_name="resize_model2")
+    model = ModelWrapper(model)
+    model.set_tensor_datatype("inp", idt)
+    model.set_tensor_datatype("outp", idt)
+    model.set_tensor_layout("inp", DataLayout.NHWC)
+
+    model = model.transform(InferShapes())
+    model = model.transform(InferDataLayouts())
+
+    return model
+
+
+def create_transpose_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
+    ofm_dim_h = ifm_dim[0] * scales[2]
+    ofm_dim_w = ifm_dim[1] * scales[3]
+    inp = oh.make_tensor_value_info(
+        "inp", TensorProto.FLOAT, [1, ifm_dim[0], ifm_dim[1], ifm_ch]
+    )
+
+    param = oh.make_tensor_value_info("scales", TensorProto.FLOAT, scales)
+
+    # Not actually used, only needed for compliance with the Resize node interface
+    roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])
+
+    outp_tr = oh.make_tensor_value_info(
+        "outp_tr", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]]
+    )
+
+    outp_up = oh.make_tensor_value_info(
+        "outp_up", TensorProto.FLOAT, [1, ifm_ch, ofm_dim_h, ofm_dim_w]
+    )
+    outp = oh.make_tensor_value_info(
+        "outp", TensorProto.FLOAT, [1, ofm_dim_h, ofm_dim_w, ifm_ch]
+    )
+
+    transpose_node1 = onnx.helper.make_node(
+        "Transpose",
+        inputs=["inp"],
+        outputs=["outp_tr"],
+        name="Transpose1",
+        perm=[0, 3, 1, 2],
+    )
+
+    resize_node = oh.make_node(
+        "Resize",
+        inputs=["outp_tr", "roi", "scales"],
+        outputs=["outp_up"],
+        name="Resize1",
+        mode=mode,
+    )
+
+    transpose_node2 = onnx.helper.make_node(
+        "Transpose",
+        inputs=["outp_up"],
+        outputs=["outp"],
+        name="Transpose2",
+        perm=[0, 2, 3, 1],
+    )
+
+    graph = oh.make_graph(
+        nodes=[transpose_node1, resize_node, transpose_node2],
+        name="resize_graph",
+        inputs=[inp],
+        outputs=[outp],
+        value_info=[outp_up, outp_tr, param, roi],
+    )
+
+    model = oh.make_model(graph, producer_name="resize_model3")
+    model = ModelWrapper(model)
+    model.set_tensor_datatype("inp", idt)
+    model.set_tensor_datatype("outp", idt)
+    model.set_tensor_layout("inp", DataLayout.NHWC)
+
+    model = model.transform(InferShapes())
+    model = model.transform(InferDataLayouts())
+
+    return model
+
+
+def check_transform(model):
+    graph = model.graph
+    node_ind = 0
+    for n in graph.node:
+        node_ind += 1
+        if n.op_type == "Upsample" or n.op_type == "Resize":
+            if model.get_tensor_layout(n.output[0]) == DataLayout.NHWC:
+                return True
+    return False
+
+
+@pytest.mark.streamline
+# input dimension
+@pytest.mark.parametrize("ifm_dim", [[2**i, 2**i] for i in range(3, 6)])
+# input channels
+@pytest.mark.parametrize("ifm_ch", [3])
+# scales
+@pytest.mark.parametrize(
+    "scales", [[1, 1, i, j] for i in range(2, 5) for j in range(2, 5)]
+)
+# mode
+@pytest.mark.parametrize("mode", ["nearest"])
+# input datatype
+@pytest.mark.parametrize("idt", [DataType["INT4"]])
+def test_scale_resize_nhwc(ifm_dim, ifm_ch, scales, mode, idt):
+    # create models
+    resize_model1 = create_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt)
+    resize_model2 = create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt)
+    resize_model3 = create_transpose_resize_transpose(
+        ifm_dim, ifm_ch, scales, mode, idt
+    )
+
+    # set initializers
+    resize_model1.set_initializer("scales", np.array(scales, dtype=np.float32))
+    resize_model2.set_initializer("scales", np.array(scales, dtype=np.float32))
+    resize_model3.set_initializer("scales", np.array(scales, dtype=np.float32))
+
+    # generate input tensor for testing
+    input_tensor_nchw = gen_finn_dt_tensor(idt, [1, ifm_ch, ifm_dim[0], ifm_dim[1]])
+    input_tensor_nhwc = gen_finn_dt_tensor(idt, [1, ifm_dim[0], ifm_dim[1], ifm_ch])
+    input_dict_nchw = {"inp": input_tensor_nchw}
+    input_dict_nhwc = {"inp": input_tensor_nhwc}
+
+    # execute first model
+    output_dict1 = oxe.execute_onnx(resize_model1, input_dict_nchw)
+    expected1 = output_dict1["outp"]
+
+    # transform Resize into ResizeNHWC
+    resize_model1 = resize_model1.transform(MakeScaleResizeNHWC())
+    resize_model1 = resize_model1.transform(InferDataLayouts())
+
+    # execute transformed model
+    output_node_name1 = resize_model1.graph.output[0].name
+    output_dict1 = oxe.execute_onnx(
+        resize_model1, input_dict_nchw, return_full_exec_context=False
+    )
+    output1 = output_dict1[output_node_name1]
+
+    # compare outputs
+    assert (expected1 == output1).all()
+    assert check_transform(resize_model1)
+
+    # execute second model
+    output_dict2 = oxe.execute_onnx(resize_model2, input_dict_nhwc)
+    expected2 = output_dict2["outp"]
+
+    # transform Resize into ResizeNHWC
+    resize_model2 = resize_model2.transform(MakeScaleResizeNHWC())
+    resize_model2 = resize_model2.transform(InferDataLayouts())
+
+    # execute transformed model
+    output_node_name2 = resize_model2.graph.output[0].name
+    output_dict2 = oxe.execute_onnx(
+        resize_model2, input_dict_nhwc, return_full_exec_context=False
+    )
+    output2 = output_dict2[output_node_name2]
+
+    # compare outputs
+    assert (expected2 == output2).all()
+    assert check_transform(resize_model2)
+
+    # execute third model
+    output_dict3 = oxe.execute_onnx(resize_model3, input_dict_nhwc)
+    expected3 = output_dict3["outp"]
+
+    # transform Resize into ResizeNHWC
+    resize_model3 = resize_model3.transform(MakeScaleResizeNHWC())
+    resize_model3 = resize_model3.transform(InferDataLayouts())
+
+    # execute transformed model
+    output_node_name3 = resize_model3.graph.output[0].name
+    output_dict3 = oxe.execute_onnx(
+        resize_model3, input_dict_nhwc, return_full_exec_context=False
+    )
+    output3 = output_dict3[output_node_name3]
+
+    # compare outputs
+    assert (expected3 == output3).all()
+    assert check_transform(resize_model3)