Extend the 4d conversion test to include resize

Author: klassen9

tests/transformation/test_4d_conversion.py

@@ -327,8 +327,67 @@ def create_conv_upsample():
             model.set_initializer(tensor_name, gen_finn_dt_tensor(DataType["FLOAT32"], init_shape))
     return model
 
+def create_resize():
+    """
+    Creates an model for testing the 3D to 4D transform of the resize node.
+    """
+    resize_node1 = onnx.helper.make_node(
+        "Resize",
+        inputs=["in_resize1", "roi_resize1", "scales_resize1", "sizes_resize1"],
+        outputs=["out_resize1"],
+        name="Resize1",
+        mode="nearest",
+    )
+
+    resize_node2 = onnx.helper.make_node(
+        "Resize",
+        inputs=["out_resize1", "roi_resize2", "scales_resize2"],
+        outputs=["out_resize2"],
+        name="Resize2",
+        mode="nearest",
+    )
+    
+    in_resize1 = onnx.helper.make_tensor_value_info("in_resize1", onnx.TensorProto.FLOAT, [1, 32, 4])
+    out_resize1 = onnx.helper.make_tensor_value_info("out_resize1", onnx.TensorProto.FLOAT, [1, 32, 8]) 
+    out_resize2 = onnx.helper.make_tensor_value_info("out_resize2", onnx.TensorProto.FLOAT, [1, 32, 16]) 
+    
+    roi_resize1 = onnx.helper.make_tensor_value_info("roi_resize1", onnx.TensorProto.FLOAT, [4])
+    scales_resize1 = onnx.helper.make_tensor_value_info("scales_resize1", onnx.TensorProto.FLOAT, [])
+    sizes_resize1 = onnx.helper.make_tensor_value_info("sizes_resize1", onnx.TensorProto.INT64, [3])
+
+    roi_resize2 = onnx.helper.make_tensor_value_info("roi_resize2", onnx.TensorProto.FLOAT, [4])
+    scales_resize2 = onnx.helper.make_tensor_value_info("scales_resize2", onnx.TensorProto.FLOAT, [3])
+    
+    list_of_nodes = [
+        resize_node1,
+        resize_node2,
+    ]
+    list_of_value_infos = [
+        out_resize1,
+        roi_resize1,
+        sizes_resize1,
+        scales_resize1,
+        roi_resize2,
+        scales_resize2,
+    ]
+
+    graph = onnx.helper.make_graph(
+        nodes=list_of_nodes,
+        name="4d_conversion_resize_test_graph",
+        inputs=[in_resize1],
+        outputs=[out_resize2],
+        value_info=list_of_value_infos,
+    )
+
+    onnx_model = qonnx_make_model(graph, producer_name="4d_conversion_resize_test-model")
+    model = ModelWrapper(onnx_model)
+    model = model.transform(InferShapes())
+    model.set_initializer("sizes_resize1", np.array([1, 32, 8], dtype=np.int64))
+    model.set_initializer("scales_resize1", np.array([], dtype=np.float32))
+    model.set_initializer("scales_resize2", np.array([1., 1., 2.], dtype=np.float32))
+    return model
 
-@pytest.mark.parametrize("test_model", ["Quartz", "VGG", "ConvUpsample"])
+@pytest.mark.parametrize("test_model", ["Quartz", "VGG", "ConvUpsample", "Resize"])
 def test_4d_conversion(test_model):
     """
     Test for the 3D to 4D transformation with a valid graph.
@@ -340,6 +399,8 @@ def test_4d_conversion(test_model):
         model = create_arbitrary_model_vgg()
     elif test_model == "ConvUpsample":
         model = create_conv_upsample()
+    elif test_model == "Resize":
+        model = create_resize()
     else:
         raise Exception("Unknown test_model in test_4d_conversion")