diff --git a/python/tvm/contrib/target/onnx.py b/python/tvm/contrib/target/onnx.py index e442c806a2f3..b839af669fe6 100644 --- a/python/tvm/contrib/target/onnx.py +++ b/python/tvm/contrib/target/onnx.py @@ -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, @@ -701,6 +791,7 @@ def convert_attributes(cls, attrs): "copy": rename("Identity"), "round": rename("Round"), "cast": Cast, + "image.resize2d": Resize, } diff --git a/tests/python/contrib/test_onnx.py b/tests/python/contrib/test_onnx.py index 8567f2c814cf..1c2d00aed866 100644 --- a/tests/python/contrib/test_onnx.py +++ b/tests/python/contrib/test_onnx.py @@ -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() @@ -684,3 +728,4 @@ def verify_cast(dshape, dtype): test_copy() test_round() test_cast() + test_resize()