You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/02/19 00:20:42 UTC
[tvm] 16/42: [Unity] Relax op: image (#13994)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit af46f0df8563d3c1548737a17a4fb247301a17f6
Author: Ruihang Lai <ru...@cs.cmu.edu>
AuthorDate: Tue Feb 14 14:57:05 2023 -0500
[Unity] Relax op: image (#13994)
This PR is about the high-level tensor computation operators in Relax.
This PR includes the image operators.
---
include/tvm/relax/attrs/image.h | 81 +++++++
python/tvm/relax/op/__init__.py | 1 +
python/tvm/relax/op/{ => image}/__init__.py | 16 +-
.../relax/op/{__init__.py => image/_ffi_api.py} | 15 +-
python/tvm/relax/op/image/image.py | 128 +++++++++++
python/tvm/relax/op/op_attrs.py | 5 +
python/tvm/script/ir_builder/relax/ir.py | 2 +
src/relax/op/image/resize.cc | 113 ++++++++++
src/relax/op/image/resize.h | 43 ++++
tests/python/relax/test_op_image.py | 245 +++++++++++++++++++++
.../python/relax/test_tvmscript_parser_op_image.py | 54 +++++
11 files changed, 678 insertions(+), 25 deletions(-)
diff --git a/include/tvm/relax/attrs/image.h b/include/tvm/relax/attrs/image.h
new file mode 100644
index 0000000000..13463aaa48
--- /dev/null
+++ b/include/tvm/relax/attrs/image.h
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relax/attrs/image.h
+ * \brief Attributes for image operators.
+ */
+#ifndef TVM_RELAX_ATTRS_IMAGE_H_
+#define TVM_RELAX_ATTRS_IMAGE_H_
+
+#include <tvm/relax/expr.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in image resize2d operator */
+struct Resize2DAttrs : public tvm::AttrsNode<Resize2DAttrs> {
+ Array<FloatImm> roi;
+ String layout;
+ String method;
+ String coordinate_transformation_mode;
+ String rounding_method;
+ double cubic_alpha;
+ int cubic_exclude;
+ double extrapolation_value;
+ DataType out_dtype;
+
+ TVM_DECLARE_ATTRS(Resize2DAttrs, "relax.attrs.Resize2DAttrs") {
+ TVM_ATTR_FIELD(roi).describe(
+ "Region of Interest for coordinate transformation mode 'tf_crop_and_resize'");
+ TVM_ATTR_FIELD(layout).describe(
+ "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc."
+ "'N', 'C', 'H', 'W' stands for batch, channel, height, and width"
+ "dimensions respectively. Resize is applied on the 'H' and"
+ "'W' dimensions.");
+ TVM_ATTR_FIELD(method).describe(
+ "Specify the mode to use for scaling."
+ "nearest_neighbor - Nearest Neighbor"
+ "linear - Bilinear Interpolation"
+ "cubic - Bicubic Interpolation");
+ TVM_ATTR_FIELD(coordinate_transformation_mode)
+ .describe(
+ "Describes how to transform the coordinate in the resized tensor"
+ "to the coordinate in the original tensor."
+ "Refer to the ONNX Resize operator specification for details"
+ "Available options are half_pixel, align_corners and asymmetric");
+ TVM_ATTR_FIELD(rounding_method)
+ .describe(
+ "indicates how to find the \"nearest\" pixel in nearest_neighbor method"
+ "Available options are round, floor, and ceil.");
+ TVM_ATTR_FIELD(cubic_alpha).describe("Spline Coefficient for Bicubic Interpolation");
+ TVM_ATTR_FIELD(cubic_exclude)
+ .describe("Flag to exclude exterior of the image during bicubic interpolation");
+ TVM_ATTR_FIELD(extrapolation_value)
+ .describe("Value to return when roi is outside of the image");
+ TVM_ATTR_FIELD(out_dtype).describe(
+ "The dtype of the output tensor. It it is not specified, the output will have the same "
+ "dtype as input if not specified.");
+ }
+}; // struct Resize2dAttrs
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_ATTRS_IMAGE_H_
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index da29c3715d..3857351269 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -26,4 +26,5 @@ from .manipulate import *
from .op_attrs import *
from .set import *
from . import builtin
+from . import image
from . import memory
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/image/__init__.py
similarity index 72%
copy from python/tvm/relax/op/__init__.py
copy to python/tvm/relax/op/image/__init__.py
index da29c3715d..f2552ad6ac 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/image/__init__.py
@@ -14,16 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
-"""Relax core operators."""
-
-# Operators
-from .base import *
-from .binary import *
-from .datatype import *
-from .index import *
-from .manipulate import *
-from .op_attrs import *
-from .set import *
-from . import builtin
-from . import memory
+# pylint: disable=wildcard-import
+"""Image operators."""
+from .image import *
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/image/_ffi_api.py
similarity index 72%
copy from python/tvm/relax/op/__init__.py
copy to python/tvm/relax/op/image/_ffi_api.py
index da29c3715d..e666203ae7 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/image/_ffi_api.py
@@ -14,16 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
-"""Relax core operators."""
+"""Constructor APIs"""
+import tvm._ffi
-# Operators
-from .base import *
-from .binary import *
-from .datatype import *
-from .index import *
-from .manipulate import *
-from .op_attrs import *
-from .set import *
-from . import builtin
-from . import memory
+tvm._ffi._init_api("relax.op.image", __name__)
diff --git a/python/tvm/relax/op/image/image.py b/python/tvm/relax/op/image/image.py
new file mode 100644
index 0000000000..562de5021d
--- /dev/null
+++ b/python/tvm/relax/op/image/image.py
@@ -0,0 +1,128 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Image operators."""
+from typing import Optional, Tuple, Union
+
+from tvm import DataType
+from tvm.ir.expr import PrimExpr
+
+from . import _ffi_api
+from ...expr import Expr, ShapeExpr
+
+
+PrimExprLike = Union[int, PrimExpr]
+
+
+def resize2d(
+ data: Expr,
+ size: Union[Expr, PrimExprLike, Tuple[PrimExprLike]],
+ roi: Optional[Union[float, Tuple[float]]] = None,
+ layout: str = "NCHW",
+ method: str = "linear",
+ coordinate_transformation_mode: str = "half_pixel",
+ rounding_method: str = "round",
+ cubic_alpha: float = -0.5,
+ cubic_exclude: int = 0,
+ extrapolation_value: float = 0.0,
+ out_dtype: Optional[Union[str, DataType]] = None,
+) -> Expr:
+ """Image resize2d operator.
+
+ This operator takes data as input and does 2D scaling to the given scale factor.
+ In the default case, where the data_layout is `NCHW`
+ with data of shape (n, c, h, w)
+ out will have a shape (n, c, size[0], size[1])
+
+ method indicates the algorithm to be used while calculating the out value
+ and method can be one of ("linear", "nearest_neighbor", "cubic")
+
+ Parameters
+ ----------
+ data : relax.Expr
+ The input data to the operator.
+
+ size: Union[Expr, PrimExprLike, Tuple[PrimExprLike]]
+ The out size to which the image will be resized.
+ If specified as a list, it is required to have length either 1 or 2.
+ If specified as an Expr, it is required to have ndim 2.
+
+ roi: Optional[Union[float, Tuple[float]]]
+ The region of interest for cropping the input image. Expected to be of
+ size 4, and format [start_h, start_w, end_h, end_w].
+ Only used if coordinate_transformation_mode is tf_crop_and_resize.
+
+ layout : str
+ Layout of the input.
+
+ method : str
+ Scale method to used [nearest_neighbor, linear, cubic].
+
+ coordinate_transformation_mode : str
+ Describes how to transform the coordinate in the resized tensor
+ to the coordinate in the original tensor. Definitions can be found
+ in topi/image/resize.py.
+ [half_pixel, align_corners, asymmetric, pytorch_half_pixel,
+ tf_half_pixel_for_nn, and tf_crop_and_resize].
+
+ rounding_method: str
+ indicates how to find the "nearest" pixel in nearest_neighbor method
+ [round, floor, ceil]
+
+ cubic_alpha: float
+ Spline Coefficient for bicubic interpolation
+
+ cubic_exclude: int
+ Flag to exclude exterior of the image during bicubic interpolation
+
+ extrapolation_value: float
+ Fill value to use when roi is outside of the image
+
+ out_dtype : Optional[Union[str, DataType]]
+ The dtype of the output tensor.
+ It it is not specified, the output will have the same dtype as input if not specified.
+
+ Returns
+ -------
+ result: relax.Expr
+ The resized result.
+ """
+ if roi is None:
+ roi = (0.0, 0.0, 0.0, 0.0) # type: ignore
+ elif isinstance(roi, float):
+ roi = (roi, roi, roi, roi) # type: ignore
+
+ if isinstance(size, (int, PrimExpr)):
+ size = (size, size)
+ if isinstance(size, tuple):
+ if len(size) == 1:
+ size = ShapeExpr([size[0], size[0]])
+ else:
+ size = ShapeExpr(size)
+
+ return _ffi_api.resize2d( # type: ignore
+ data,
+ size,
+ roi,
+ layout,
+ method,
+ coordinate_transformation_mode,
+ rounding_method,
+ cubic_alpha,
+ cubic_exclude,
+ extrapolation_value,
+ out_dtype,
+ )
diff --git a/python/tvm/relax/op/op_attrs.py b/python/tvm/relax/op/op_attrs.py
index 47c3b28798..fb64443b7e 100644
--- a/python/tvm/relax/op/op_attrs.py
+++ b/python/tvm/relax/op/op_attrs.py
@@ -34,6 +34,11 @@ class StridedSliceAttrs(Attrs):
"""Attributes used in strided_slice operator"""
+@tvm._ffi.register_object("relax.attrs.Resize2DAttrs")
+class Resize2DAttrs(Attrs):
+ """Attributes used in image resize2d operator"""
+
+
@tvm._ffi.register_object("relax.attrs.UniqueAttrs")
class UniqueAttrs(Attrs):
"""Attributes used for the unique operator"""
diff --git a/python/tvm/script/ir_builder/relax/ir.py b/python/tvm/script/ir_builder/relax/ir.py
index 537adec615..22b85f6f40 100644
--- a/python/tvm/script/ir_builder/relax/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -35,6 +35,7 @@ from tvm.relax.op import (
builtin,
call_builtin_with_ctx,
call_tir,
+ image,
invoke_closure,
make_closure,
memory,
@@ -420,6 +421,7 @@ __all__ = [
"func_ret_struct_info",
"func_ret_value",
"function",
+ "image",
"invoke_closure",
"make_closure",
"memory",
diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc
new file mode 100644
index 0000000000..2711b3cc45
--- /dev/null
+++ b/src/relax/op/image/resize.cc
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file resize.cc
+ * \brief Image resize operators.
+ */
+
+#include "resize.h"
+
+#include <utility>
+
+namespace tvm {
+namespace relax {
+
+/* relax.resize2d */
+TVM_REGISTER_NODE_TYPE(Resize2DAttrs);
+
+Expr resize2d(Expr data, Expr size, Array<FloatImm> roi, String layout, String method,
+ String coordinate_transformation_mode, String rounding_method, double cubic_alpha,
+ int cubic_exclude, double extrapolation_value, DataType out_dtype) {
+ ObjectPtr<Resize2DAttrs> attrs = make_object<Resize2DAttrs>();
+ attrs->roi = std::move(roi);
+ attrs->layout = std::move(layout);
+ attrs->method = std::move(method);
+ attrs->coordinate_transformation_mode = std::move(coordinate_transformation_mode);
+ attrs->rounding_method = std::move(rounding_method);
+ attrs->cubic_alpha = cubic_alpha;
+ attrs->cubic_exclude = cubic_exclude;
+ attrs->extrapolation_value = extrapolation_value;
+ attrs->out_dtype = out_dtype;
+
+ static const Op& op = Op::Get("relax.image.resize2d");
+ return Call(op, {std::move(data), std::move(size)}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.image.resize2d").set_body_typed(resize2d);
+
+StructInfo InferStructInfoResize2D(const Call& call, const BlockBuilder& ctx) {
+ if (call->args.size() != 1 && call->args.size() != 2) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "Resize2D expects either one or two arguments, while the given number of arguments is "
+ << call->args.size());
+ }
+
+ const auto* data_sinfo = GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
+ const auto* size_sinfo = GetStructInfoAs<ShapeStructInfoNode>(call->args[1]);
+ const auto* size_value = call->args[1].as<ShapeExprNode>();
+ if (data_sinfo == nullptr) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "Resize2D expects the input data to be a Tensor, while the given data is "
+ << call->args[0]->GetTypeKey());
+ }
+ if (size_sinfo == nullptr) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "Resize2D expects the given output image size to be a Shape, while the given one is "
+ << call->args[1]->GetTypeKey());
+ }
+ if (size_sinfo->ndim != 2) {
+ ctx->ReportFatal(Diagnostic::Error(call) << "Resize2D expects the given output image size to "
+ "be a 2-dim shape, while the given one has ndim "
+ << size_sinfo->ndim);
+ }
+
+ const auto* attrs = call->attrs.as<Resize2DAttrs>();
+ auto [data_layout, data2NCHW] = CheckTensorLayout(call, ctx, attrs->layout, //
+ /*tgt_layout=*/"NCHW", //
+ /*tensor_name=*/"data");
+
+ DataType out_dtype = attrs->out_dtype.is_void() ? data_sinfo->dtype : attrs->out_dtype;
+
+ Optional<ShapeExpr> data_shape =
+ CheckNdimPerLayoutAndGetShape(call, ctx, GetRef<TensorStructInfo>(data_sinfo), data_layout);
+ if (!data_shape.defined() || size_value == nullptr) {
+ return TensorStructInfo(out_dtype, data_layout.ndim());
+ }
+
+ Array<PrimExpr> data_NCHW_shape = data2NCHW.ForwardShape(data_shape.value()->values);
+ Array<PrimExpr> out_NCHW_shape(data_NCHW_shape);
+ out_NCHW_shape.Set(2, size_value->values[0]);
+ out_NCHW_shape.Set(3, size_value->values[1]);
+
+ Array<PrimExpr> out_shape = data2NCHW.BackwardShape(out_NCHW_shape);
+ return TensorStructInfo(ShapeExpr(out_shape), out_dtype);
+}
+
+TVM_REGISTER_OP("relax.image.resize2d")
+ .set_attrs_type<Resize2DAttrs>()
+ .set_num_inputs(2)
+ .add_argument("data", "Tensor", "The input tensor.")
+ .add_argument("size", "Shape", "The output image shape.")
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoResize2D);
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/op/image/resize.h b/src/relax/op/image/resize.h
new file mode 100644
index 0000000000..085a1cbc5d
--- /dev/null
+++ b/src/relax/op/image/resize.h
@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file resize.h
+ * \brief The functions to make Relax image resize operator calls.
+ */
+
+#ifndef TVM_RELAX_OP_IMAGE_RESIZE_H_
+#define TVM_RELAX_OP_IMAGE_RESIZE_H_
+
+#include <tvm/relax/attrs/image.h>
+
+#include "../op_common.h"
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Image resize2d operator. */
+Expr resize2d(Expr data, Expr size, Array<FloatImm> roi, String layout, String method,
+ String coordinate_transformation_mode, String rounding_method, double cubic_alpha,
+ int cubic_exclude, double extrapolation_value, DataType out_dtype);
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_OP_IMAGE_RESIZE_H_
diff --git a/tests/python/relax/test_op_image.py b/tests/python/relax/test_op_image.py
new file mode 100644
index 0000000000..b06b51a2a1
--- /dev/null
+++ b/tests/python/relax/test_op_image.py
@@ -0,0 +1,245 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pytest
+import tvm
+import tvm.testing
+from tvm import relax, tir
+from tvm import TVMError
+from tvm.ir import Op
+from tvm.script import relax as R
+
+
+def test_op_correctness():
+ x = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+ assert relax.op.image.resize2d(x, (28, 28)).op == Op.get("relax.image.resize2d")
+
+
+def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo: relax.StructInfo):
+ ret = bb.normalize(call)
+ tvm.ir.assert_structural_equal(ret.struct_info, expected_sinfo)
+
+
+def test_resize2d_infer_struct_info():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+ x1 = relax.Var("x", R.Tensor((2, 32, 32, 3), "float32"))
+ x2 = relax.Var("x", R.Tensor((2, 4, 32, 32, 16), "float32"))
+ x3 = relax.Var("x", R.Tensor("float32", ndim=4))
+ x4 = relax.Var("x", R.Tensor("float32", ndim=5))
+ x5 = relax.Var("x", R.Tensor("float32"))
+ x6 = relax.Var("x", R.Tensor(ndim=4))
+ x7 = relax.Var("x", R.Tensor())
+
+ _check_inference(
+ bb, relax.op.image.resize2d(x0, (28, 28)), relax.TensorStructInfo((2, 3, 28, 28), "float32")
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x0, size=28),
+ relax.TensorStructInfo((2, 3, 28, 28), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x0, size=(28, 30)),
+ relax.TensorStructInfo((2, 3, 28, 30), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x1, size=28, layout="NHWC"),
+ relax.TensorStructInfo((2, 28, 28, 3), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x0, size=28, out_dtype="float16"),
+ relax.TensorStructInfo((2, 3, 28, 28), "float16"),
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x2, size=28, layout="NCHW16c"),
+ relax.TensorStructInfo((2, 4, 28, 28, 16), "float32"),
+ )
+ _check_inference(
+ bb, relax.op.image.resize2d(x3, size=28), relax.TensorStructInfo(dtype="float32", ndim=4)
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x4, size=28, layout="NCHW16c"),
+ relax.TensorStructInfo(dtype="float32", ndim=5),
+ )
+ _check_inference(
+ bb, relax.op.image.resize2d(x5, size=28), relax.TensorStructInfo(dtype="float32", ndim=4)
+ )
+ _check_inference(
+ bb, relax.op.image.resize2d(x6, size=28), relax.TensorStructInfo(dtype="", ndim=4)
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x6, size=28, out_dtype="float32"),
+ relax.TensorStructInfo(dtype="float32", ndim=4),
+ )
+ _check_inference(
+ bb, relax.op.image.resize2d(x7, size=28), relax.TensorStructInfo(dtype="", ndim=4)
+ )
+
+
+def test_resize2d_infer_struct_info_shape_symbolic():
+ bb = relax.BlockBuilder()
+ n = tir.Var("n", "int64")
+ c = tir.Var("c", "int64")
+ ih = tir.Var("ih", "int64")
+ iw = tir.Var("iw", "int64")
+ oh = tir.Var("oh", "int64")
+ ow = tir.Var("ow", "int64")
+ x0 = relax.Var("x", R.Tensor((n, c, ih, iw), "float32"))
+ x1 = relax.Var("x", R.Tensor((n, c, ih, iw, 16), "float32"))
+
+ _check_inference(
+ bb, relax.op.image.resize2d(x0, size=oh), relax.TensorStructInfo((n, c, oh, oh), "float32")
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x0, size=(oh, ow)),
+ relax.TensorStructInfo((n, c, oh, ow), "float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x1, size=(oh, ow), layout="NCHW16c"),
+ relax.TensorStructInfo((n, c, oh, ow, 16), "float32"),
+ )
+
+
+def test_resize2d_infer_struct_info_shape_var():
+ bb = relax.BlockBuilder()
+ s0 = relax.Var("s", relax.ShapeStructInfo(ndim=4))
+ s1 = relax.Var("s", relax.ShapeStructInfo(ndim=5))
+ s2 = relax.Var("s", relax.ShapeStructInfo())
+ x0 = relax.Var("x", relax.TensorStructInfo(s0, "float32"))
+ x1 = relax.Var("x", relax.TensorStructInfo(s1, "float32"))
+ x2 = relax.Var("x", relax.TensorStructInfo(s2, "float32"))
+
+ _check_inference(
+ bb, relax.op.image.resize2d(x0, size=32), relax.TensorStructInfo(dtype="float32", ndim=4)
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x1, size=32, layout="NCHW16c"),
+ relax.TensorStructInfo(dtype="float32", ndim=5),
+ )
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x2, size=32, layout="NCHW16c"),
+ relax.TensorStructInfo(dtype="float32", ndim=5),
+ )
+
+
+def test_resize2d_infer_struct_info_pool_size_var():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+ s0 = relax.Var("s", relax.ShapeStructInfo((30, 30)))
+ s1 = relax.Var("s", relax.ShapeStructInfo(ndim=2))
+
+ _check_inference(
+ bb,
+ relax.op.image.resize2d(x0, s0),
+ relax.TensorStructInfo(dtype="float32", ndim=4),
+ )
+ _check_inference(
+ bb, relax.op.image.resize2d(x0, s1), relax.TensorStructInfo(dtype="float32", ndim=4)
+ )
+
+
+def test_resize2d_infer_struct_info_more_input_dtype():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float16"))
+ x1 = relax.Var("x", R.Tensor((2, 3, 32, 32), "int8"))
+ x2 = relax.Var("x", R.Tensor((2, 3, 32, 32), "int64"))
+ _check_inference(
+ bb, relax.op.image.resize2d(x0, size=28), relax.TensorStructInfo((2, 3, 28, 28), "float16")
+ )
+ _check_inference(
+ bb, relax.op.image.resize2d(x1, size=28), relax.TensorStructInfo((2, 3, 28, 28), "int8")
+ )
+ _check_inference(
+ bb, relax.op.image.resize2d(x2, size=28), relax.TensorStructInfo((2, 3, 28, 28), "int64")
+ )
+
+
+def test_resize2d_infer_struct_info_wrong_layout_string():
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x, size=28, layout="OIHW"))
+
+
+def test_resize2d_wrong_input_ndim():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+ x1 = relax.Var("x", R.Tensor((2, 3, 32, 32, 3), "float32"))
+ x2 = relax.Var("x", R.Tensor("float32", ndim=3))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, size=28, layout="NCHW16c"))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x1, size=28, layout="NCHW"))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x2, size=28))
+
+
+def test_resize2d_wrong_pool_size_ndim():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float16"))
+ s0 = relax.ShapeExpr((3,))
+ s1 = relax.Var("s", relax.ShapeStructInfo((30, 30, 30)))
+ s2 = relax.Var("s", relax.ShapeStructInfo(ndim=3))
+ s3 = relax.Var("s", relax.ShapeStructInfo(ndim=1))
+ s4 = relax.Var("s", relax.ShapeStructInfo(ndim=0))
+ s5 = relax.Var("s", relax.ShapeStructInfo())
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, (3, 3, 3)))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, s0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, s1))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, s2))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, s3))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, s4))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, s5))
+
+
+def test_resize2d_infer_struct_info_wrong_input_type():
+ bb = relax.BlockBuilder()
+ x0 = relax.Var("x", relax.ShapeStructInfo((2, 3, 28, 28)))
+ x1 = relax.Var("x", relax.FuncStructInfo([], R.Tensor((2, 3, 28, 28), "float32")))
+ x2 = relax.Var("x", R.Tensor((2, 3, 32, 32), "float32"))
+ s0 = relax.Var("s", R.Tensor((3, 3)))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x0, size=32))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x1, size=32))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.image.resize2d(x2, s0))
+ with pytest.raises(TVMError):
+ relax.op.image.resize2d(x2, [30, 30])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_tvmscript_parser_op_image.py b/tests/python/relax/test_tvmscript_parser_op_image.py
new file mode 100644
index 0000000000..a90da37812
--- /dev/null
+++ b/tests/python/relax/test_tvmscript_parser_op_image.py
@@ -0,0 +1,54 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Optional, Union
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm import IRModule, relax
+from tvm.script import relax as R
+
+
+def _check(
+ parsed: Union[relax.Function, IRModule],
+ expect: Optional[Union[relax.Function, IRModule]],
+):
+ test = parsed.script(show_meta=True)
+ roundtrip_mod = tvm.script.from_source(test)
+ tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
+ if expect:
+ tvm.ir.assert_structural_equal(parsed, expect)
+
+
+def test_resize2d():
+ @R.function
+ def foo(x: R.Tensor((2, 14, 14, 3), "float32")) -> R.Tensor((2, 28, 28, 3), "float32"):
+ gv: R.Tensor((2, 28, 28, 3), "float32") = R.image.resize2d(x, size=(28, 28), layout="NHWC")
+ return gv
+
+ bb = relax.BlockBuilder()
+ x = relax.Var("x", R.Tensor((2, 14, 14, 3), "float32"))
+ with bb.function("foo", [x]):
+ gv = bb.emit(relax.op.image.resize2d(x, (28, 28), layout="NHWC"))
+ bb.emit_func_output(gv)
+
+ _check(foo, bb.get()["foo"])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()