Skip to content

Commit

Permalink
[Relay] Support resize in the ONNX conversion (apache#8455)
Browse files Browse the repository at this point in the history
* [Relay to Onnx]

* Added support for resize2d op
* Added unit test

* [Relay to Onnx][Resize]

* Fixed formatting errors

* [Relay to Onnx][Resize]

* Fixed issue in resize conversion: round maps to round_preferc_ceil
* Updated resize unit test to test for coordinate transform mode and
round
* Known issue: Does not match for (NN, align_corners) and Cubic

* * Fixed formatting errors

* * Fixed some more formatting errors
  • Loading branch information
schilkunda-amba authored and ylc committed Sep 29, 2021
1 parent 3e90bc3 commit 9d58481
Show file tree
Hide file tree
Showing 2 changed files with 136 additions and 0 deletions.
91 changes: 91 additions & 0 deletions python/tvm/contrib/target/onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -662,6 +662,96 @@ def convert_attributes(cls, attrs):
return {"to": getattr(TensorProto, attrs.dtype.upper())}


class Resize(OpConverter):
"""Operator converter for Resize."""

@classmethod
def convert_attributes(cls, attrs):
method = attrs.get_str("method")
if method == "nearest_neighbor":
mode = b"nearest"
elif "linear" in method: # linear / bilinear
mode = b"linear"
elif "cubic" in method: # cubic / bicubic
mode = b"cubic"
else:
raise RuntimeError("Unsupported method %s in operator Resize" % method)

coord_trans = attrs.get_str("coordinate_transformation_mode")
if coord_trans == "half_pixel":
coord_trans = b"half_pixel"
elif coord_trans == "align_corners":
coord_trans = b"align_corners"
elif coord_trans == "asymmetric":
coord_trans = b"asymmetric"
else:
raise RuntimeError(
"Unsupported coordinate transform mode %s in operator Resize" % coord_trans
)

rounding_method = attrs.get_str("rounding_method")
if rounding_method == "round":
rounding_method = b"round_prefer_ceil"
elif rounding_method == "floor":
rounding_method = b"floor"
elif rounding_method == "ceil":
rounding_method = b"ceil"
else:
raise RuntimeError(
"Unsupported rounding method %s in operator Resize" % rounding_method
)

size = attrs.get_int_tuple("size")

return {
"mode": mode,
"coord_trans": coord_trans,
"size": size,
"nearest_mode": rounding_method,
}

@classmethod
def convert(cls, node_entry, model_container, node_dict):
attrs = cls.convert_attributes(node_entry["relay_node"].attrs)

name = node_entry["name"]
input_node = node_dict[node_entry["inputs"][0]]
assert len(input_node) == 1, "input node can not be a Tuple"
input_node = input_node[0]
input_shape = input_node["types"][0].shape

# (TBD) needed in opset 11
roi = [0] * len(input_shape) + [1] * len(input_shape)
roi_array = numpy.asarray(roi).astype(numpy.float64)
roi_node = add_input(roi_array, name, "roi", model_container)

out_size = attrs["size"]

# (onnx) rank of scale / size must match rank of X
# relay size node contains only spatial dimensions
# pad with 1s to match rank
match_rank_pad = len(input_shape) - len(out_size)
out_size_full_rank = input_shape[:match_rank_pad] + list(out_size)
out_size_array = numpy.asarray(out_size_full_rank).astype(numpy.int64)

input_size_array = numpy.asarray(list(input_shape)).astype(numpy.int64)

scale_array = numpy.divide(out_size_array, input_size_array).astype(numpy.float32)
scale_node = add_input(scale_array, name, "scales", model_container)

input_names = [node_entry["input_names"][0], roi_node, scale_node]

resize_node = onnx.helper.make_node(
cls.__name__,
input_names,
node_entry["output_names"],
mode=attrs["mode"],
coordinate_transformation_mode=attrs["coord_trans"],
nearest_mode=attrs["nearest_mode"],
)
model_container.add_nodes([resize_node])


relay_to_onnx_op_mapping = {
"reshape": Reshape,
"nn.conv2d": Conv,
Expand Down Expand Up @@ -701,6 +791,7 @@ def convert_attributes(cls, attrs):
"copy": rename("Identity"),
"round": rename("Round"),
"cast": Cast,
"image.resize2d": Resize,
}


Expand Down
45 changes: 45 additions & 0 deletions tests/python/contrib/test_onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -655,6 +655,50 @@ def verify_cast(dshape, dtype):
verify_cast(i, o_dtype)


def test_resize():
"""Resize unit test."""

def verify_resize(dshape, outsize, method, coord_trans, rounding_method, dtype="float32"):
x = relay.var("x", relay.ty.TensorType(dshape, dtype))
y = relay.image.resize2d(
x,
outsize,
layout="NCHW",
method=method,
coordinate_transformation_mode=coord_trans,
rounding_method=rounding_method,
)
func = relay.Function([x], y)
x_data = np.random.uniform(size=dshape).astype(dtype)
verify_results(func, [x_data], "test_resize", rtol=1e-4, atol=1e-4)

method = ["nearest_neighbor", "linear", "cubic"]
coord_trans = ["half_pixel", "align_corners", "asymmetric"]
rounding_method = ["round", "floor", "ceil"]

isize = (1, 3, 480, 640)

# Downsample
osize = (240, 320)
for i in method:
for j in coord_trans:
for k in rounding_method:
if (i == "nearest_neighbor" and j == "align_corners") or (
i == "cubic" and j in ["half_pixel", "align_corners"]
):
continue
verify_resize(isize, osize, method=i, coord_trans=j, rounding_method=k)

# Upsample
osize = (960, 1280)
for i in method:
for j in coord_trans:
for k in rounding_method:
if (i == "nearest_neighbor" and j == "align_corners") or (i == "cubic"):
continue
verify_resize(isize, osize, method=i, coord_trans=j, rounding_method=k)


if __name__ == "__main__":
test_add()
test_bias_add()
Expand Down Expand Up @@ -684,3 +728,4 @@ def verify_cast(dshape, dtype):
test_copy()
test_round()
test_cast()
test_resize()

0 comments on commit 9d58481

Please sign in to comment.