diff --git a/src/finn/transformation/streamline/reorder.py b/src/finn/transformation/streamline/reorder.py
index 89dd2f5a616d96ce18bf1b4786b32ec985f3f910..e36de2aa54f7b260016dd0b1b93bccbfca9d72b9 100644
--- a/src/finn/transformation/streamline/reorder.py
+++ b/src/finn/transformation/streamline/reorder.py
@@ -722,11 +722,13 @@ class MakeMaxPoolNHWC(Transformation):
graph_modified = True
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
@@ -738,12 +740,15 @@ class MakeScaleResizeNHWC(Transformation):
if n.op_type == "Upsample":
scales_ind = 1
else:
- scales_ind = 2
+ 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'))
+ 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]
@@ -762,7 +767,10 @@ class MakeScaleResizeNHWC(Transformation):
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'))
+ 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]
diff --git a/tests/transformation/streamline/test_scale_resize_nhwc.py b/tests/transformation/streamline/test_scale_resize_nhwc.py
index ba4df84569a706d7f273b6e7e48974e2f1760ee3..06faa8371929a44b95cdd2259ab637262395b78d 100644
--- a/tests/transformation/streamline/test_scale_resize_nhwc.py
+++ b/tests/transformation/streamline/test_scale_resize_nhwc.py
@@ -1,7 +1,7 @@
import pytest
+
import numpy as np
import onnx
-import onnx.version_converter as vc
import onnx.helper as oh
from onnx import TensorProto
from qonnx.core.datatype import DataType
@@ -12,6 +12,7 @@ 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]
@@ -19,14 +20,10 @@ def create_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
"inp", TensorProto.FLOAT, [1, ifm_ch, ifm_dim[0], ifm_dim[1]]
)
- param = oh.make_tensor_value_info(
- "scales", TensorProto.FLOAT, [4]
- )
+ 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]
- )
+ 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]
@@ -39,7 +36,7 @@ def create_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
"Resize",
inputs=["inp", "roi", "scales"],
outputs=["outp_up"],
- name = "Resize1",
+ name="Resize1",
mode=mode,
)
@@ -68,6 +65,7 @@ def create_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
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]
@@ -75,16 +73,12 @@ def create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt):
"inp", TensorProto.FLOAT, [1, ifm_dim[0], ifm_dim[1], ifm_ch]
)
- param = oh.make_tensor_value_info(
- "scales", TensorProto.FLOAT, [4]
- )
+ 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]
- )
+ roi = oh.make_tensor_value_info("roi", TensorProto.FLOAT, [4])
- outp= oh.make_tensor_value_info(
+ 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(
@@ -103,7 +97,7 @@ def create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt):
"Resize",
inputs=["outp_tr", "roi", "scales"],
outputs=["outp"],
- name = "Resize1",
+ name="Resize1",
mode=mode,
)
@@ -114,7 +108,7 @@ def create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt):
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)
@@ -124,6 +118,7 @@ def create_transpose_resize(ifm_dim, ifm_ch, scales, mode, idt):
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]
@@ -131,14 +126,10 @@ def create_transpose_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
"inp", TensorProto.FLOAT, [1, ifm_dim[0], ifm_dim[1], ifm_ch]
)
- param = oh.make_tensor_value_info(
- "scales", TensorProto.FLOAT, scales
- )
+ 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]
- )
+ 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]]
@@ -163,7 +154,7 @@ def create_transpose_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
"Resize",
inputs=["outp_tr", "roi", "scales"],
outputs=["outp_up"],
- name = "Resize1",
+ name="Resize1",
mode=mode,
)
@@ -195,11 +186,13 @@ def create_transpose_resize_transpose(ifm_dim, ifm_ch, scales, mode, idt):
@pytest.mark.streamline
# input dimension
-@pytest.mark.parametrize("ifm_dim", [[2**i, 2**i] for i in range(3,6)])
+@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)])
+@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
@@ -208,9 +201,11 @@ 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)
+ resize_model3 = create_transpose_resize_transpose(
+ ifm_dim, ifm_ch, scales, mode, idt
+ )
- #set initializers
+ # 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))
@@ -221,7 +216,6 @@ def test_scale_resize_nhwc(ifm_dim, ifm_ch, scales, mode, idt):
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"]
@@ -263,12 +257,12 @@ def test_scale_resize_nhwc(ifm_dim, ifm_ch, scales, mode, idt):
# transform Resize into ResizeNHWC
resize_model3 = resize_model3.transform(MakeScaleResizeNHWC())
- # execute transformed model
+ # 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
+ # compare outputs
assert (expected3 == output3).all()