[Relay] Support resize in the ONNX conversion

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,
 }
 
 

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()