1827 lines
67 KiB
Python
1827 lines
67 KiB
Python
# 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, tirx
|
|
from tvm.ir import Op, VDevice
|
|
from tvm.script import relax as R
|
|
|
|
|
|
def test_conv1d_op_correctness():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
assert relax.op.nn.conv1d(x, w).op == Op.get("relax.nn.conv1d")
|
|
assert relax.op.nn.conv1d_transpose(x, w).op == Op.get("relax.nn.conv1d_transpose")
|
|
|
|
|
|
def test_conv2d_op_correctness():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
assert relax.op.nn.conv2d(x, w).op == Op.get("relax.nn.conv2d")
|
|
assert relax.op.nn.conv2d_transpose(x, w).op == Op.get("relax.nn.conv2d_transpose")
|
|
|
|
|
|
def test_conv3d_op_correctness():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3, 3, 3), "float32"))
|
|
assert relax.op.nn.conv3d(x, w).op == Op.get("relax.nn.conv3d")
|
|
wt = relax.Var("wt", R.Tensor((3, 4, 3, 3, 3), "float32"))
|
|
assert relax.op.nn.conv3d_transpose(x, wt).op == Op.get("relax.nn.conv3d_transpose")
|
|
|
|
|
|
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_ty: relax.Type):
|
|
ret = bb.normalize(call)
|
|
tvm.ir.assert_structural_equal(ret.ty, expected_ty)
|
|
|
|
|
|
def test_conv1d_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x3 = relax.Var("x", R.Tensor("float32"))
|
|
x4 = relax.Var("x", R.Tensor())
|
|
x5 = relax.Var("x", R.Tensor((2, 4, 28, 16), "float32"))
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 28), "float32", vdev0))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=3))
|
|
w3 = relax.Var("w", R.Tensor("float32"))
|
|
w4 = relax.Var("w", R.Tensor((48, 4, 3, 16), "float32"))
|
|
w5 = relax.Var("w", R.Tensor((4, 3, 3), "float32", vdev0))
|
|
|
|
_check_inference(bb, relax.op.nn.conv1d(x0, w0), relax.TensorType((2, 4, 26), "float32"))
|
|
_check_inference(bb, relax.op.nn.conv1d(x6, w5), relax.TensorType((2, 4, 26), "float32", vdev0))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, out_dtype="float16"),
|
|
relax.TensorType((2, 4, 26), "float16"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d(x0, w0, padding=1), relax.TensorType((2, 4, 28), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, padding=[1, 3]),
|
|
relax.TensorType((2, 4, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, strides=2),
|
|
relax.TensorType((2, 4, 13), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, strides=(2,)),
|
|
relax.TensorType((2, 4, 13), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, dilation=2),
|
|
relax.TensorType((2, 4, 24), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, dilation=(2,)),
|
|
relax.TensorType((2, 4, 24), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x1, w0, data_layout="NWC"),
|
|
relax.TensorType((2, 26, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, out_layout="NWC"),
|
|
relax.TensorType((2, 26, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w1, kernel_layout="IOW"),
|
|
relax.TensorType((2, 4, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(
|
|
x5, w4, data_layout="NCW16c", kernel_layout="OIW16i", out_layout="NWC16c"
|
|
),
|
|
relax.TensorType((2, 26, 3, 16), "float32"),
|
|
)
|
|
_check_inference(bb, relax.op.nn.conv1d(x2, w0), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(bb, relax.op.nn.conv1d(x3, w0), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(bb, relax.op.nn.conv1d(x0, w2), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(bb, relax.op.nn.conv1d(x0, w3), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(bb, relax.op.nn.conv1d(x4, w0), relax.TensorType(dtype="", ndim=3))
|
|
|
|
|
|
def test_conv1d_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
c16 = tirx.Var("c16", "int64")
|
|
iw = tirx.Var("iw", "int64")
|
|
ki = tirx.Var("ki", "int64")
|
|
ko = tirx.Var("ko", "int64")
|
|
kw = tirx.Var("kw", "int64")
|
|
x0 = relax.Var("x", R.Tensor((n, c, iw), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, c, iw, c16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((ko, ki, kw), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((ko, c, kw), "float32"))
|
|
w2 = relax.Var("w", R.Tensor((ko, c, kw, c16), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0),
|
|
relax.TensorType((n, ko, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w1),
|
|
relax.TensorType((n, ko, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x1, w2, data_layout="NCW16c", kernel_layout="OIW16i", out_layout="NCW"),
|
|
relax.TensorType((n, ko, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, strides=2, padding=1, dilation=2),
|
|
relax.TensorType(
|
|
(n, ko, tvm.tirx.floordiv(iw + 3, 2) + 1 - kw),
|
|
"float32",
|
|
),
|
|
)
|
|
|
|
|
|
def test_conv1d_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s2 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s3 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
w = relax.Var("w", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(bb, relax.op.nn.conv1d(x0, w), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x1, w, data_layout="NCW16c"),
|
|
relax.TensorType(dtype="float32", ndim=4),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w, out_layout="NCW16c"),
|
|
relax.TensorType(dtype="float32", ndim=4),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x2, w),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
|
|
|
|
def test_conv1d_infer_ty_groups():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 8, 28, 16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((48, 16, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((48, 2, 3, 8), "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d(x0, w0, groups=8), relax.TensorType((2, 48, 26), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w1, kernel_layout="OIW8i", groups=8),
|
|
relax.TensorType((2, 48, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x1, w0, data_layout="NCW16c", groups=8),
|
|
relax.TensorType((2, 3, 26, 16), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv1d_infer_ty_symbolic_groups():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x = relax.Var("x", R.Tensor((n, ic * 4, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((oc * 4, ic, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((oc, ic, 3), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x, w0, groups=4),
|
|
relax.TensorType((n, oc * 4, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d(x, w1, groups=4), relax.TensorType((n, oc, 26), "float32")
|
|
)
|
|
|
|
|
|
def test_conv1d_infer_ty_input_channel_group_incompatible():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((48, 20, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, ic * 6, 28), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((oc, ic - 1, 3), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x0, w0, groups=6))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x1, w1, groups=6))
|
|
|
|
|
|
def test_conv1d_infer_ty_output_channel_group_incompatible():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x0 = relax.Var("x", R.Tensor((2, 120, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((128, 20, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, ic * 6, 28), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((oc * 6 + 4, ic * 6, 3), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x0, w0, groups=6))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x1, w1, groups=6))
|
|
|
|
|
|
def test_conv1d_non_positive_group():
|
|
x = relax.Var("x", R.Tensor((2, 128, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((48, 16, 3), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d(x, w, groups=0)
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d(x, w, groups=-2)
|
|
|
|
|
|
def test_conv1d_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28), "float64"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 3), "float64"))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 28), "int8"))
|
|
w2 = relax.Var("w", R.Tensor((4, 3, 3), "int8"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 28), "int32"))
|
|
w3 = relax.Var("w", R.Tensor((4, 3, 3), "int32"))
|
|
|
|
_check_inference(bb, relax.op.nn.conv1d(x0, w0), relax.TensorType((2, 4, 26), "float16"))
|
|
_check_inference(bb, relax.op.nn.conv1d(x1, w1), relax.TensorType((2, 4, 26), "float64"))
|
|
_check_inference(bb, relax.op.nn.conv1d(x2, w2), relax.TensorType((2, 4, 26), "int8"))
|
|
_check_inference(bb, relax.op.nn.conv1d(x3, w3), relax.TensorType((2, 4, 26), "int32"))
|
|
|
|
|
|
def test_conv1d_infer_ty_mixed_precision():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28), "int8"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 3), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 28)))
|
|
w2 = relax.Var("w", R.Tensor((4, 3, 3)))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w0, out_dtype="float32"),
|
|
relax.TensorType((2, 4, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x1, w1, out_dtype="int32"),
|
|
relax.TensorType((2, 4, 26), "int32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x2, w2, out_dtype="float32"),
|
|
relax.TensorType((2, 4, 26), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv1d_unequal_input_channel():
|
|
bb = relax.BlockBuilder()
|
|
ic = tirx.Var("ic", "int64")
|
|
x0 = relax.Var("x", R.Tensor([2, 3, 28], "float32"))
|
|
w0 = relax.Var("w", R.Tensor([3, 4, 3], "float32"))
|
|
x1 = relax.Var("x", R.Tensor([2, ic, 28], "float32"))
|
|
w1 = relax.Var("w", R.Tensor([4, ic + 2, 3], "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x0, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x1, w1))
|
|
|
|
|
|
def test_conv1d_stride_padding_dilation_int64():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
conv1d = relax.op.nn.conv1d(x, w, strides=(1,), padding=(1, 1), dilation=(1,))
|
|
|
|
assert isinstance(conv1d.attrs.strides[0], int)
|
|
assert isinstance(conv1d.attrs.padding[0], int)
|
|
assert isinstance(conv1d.attrs.padding[1], int)
|
|
assert isinstance(conv1d.attrs.dilation[0], int)
|
|
|
|
|
|
def test_conv1d_wrong_strides_padding_dilation_length():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d(x, w, strides=(1, 2))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d(x, w, padding=(1, 2, 3))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d(x, w, dilation=(1, 2))
|
|
|
|
|
|
def test_conv1d_infer_ty_wrong_layout_string():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x, w, data_layout="OIW"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x, w, kernel_layout="NWC"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x, w, out_layout="OWI"))
|
|
|
|
|
|
def test_conv1d_dtype_mismatch():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3), "int8"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv1d(x, w))
|
|
|
|
|
|
def test_conv1d_wrong_input_ndim():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 6, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=5))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x0, w1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x0, w1, data_layout="NCW16c"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x0, w2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x1, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d(x2, w0))
|
|
|
|
|
|
def test_conv1d_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
x1 = relax.Var("x", relax.ShapeType((2, 3, 28)))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
w1 = relax.Var("w", relax.FuncType([], R.Tensor((4, 3, 3), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv1d(x0, w1))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv1d(x1, w0))
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
x3 = relax.Var("x", R.Tensor("float32"))
|
|
x4 = relax.Var("x", R.Tensor())
|
|
x5 = relax.Var("x", R.Tensor((2, 4, 28, 16), "float32"))
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 28), "float32", vdev0))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=3))
|
|
w3 = relax.Var("w", R.Tensor("float32"))
|
|
w4 = relax.Var("w", R.Tensor((4, 48, 3, 16), "float32"))
|
|
w5 = relax.Var("w", R.Tensor((3, 4, 3), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x0, w0), relax.TensorType((2, 4, 30), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x6, w5),
|
|
relax.TensorType((2, 4, 30), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, out_dtype="float16"),
|
|
relax.TensorType((2, 4, 30), "float16"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, padding=1),
|
|
relax.TensorType((2, 4, 28), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, padding=[1, 3]),
|
|
relax.TensorType((2, 4, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, strides=3, output_padding=1),
|
|
relax.TensorType((2, 4, 85), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, strides=2),
|
|
relax.TensorType((2, 4, 57), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, dilation=2),
|
|
relax.TensorType((2, 4, 32), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, dilation=(2,)),
|
|
relax.TensorType((2, 4, 32), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x1, w0, data_layout="NWC"),
|
|
relax.TensorType((2, 30, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, out_layout="NWC"),
|
|
relax.TensorType((2, 30, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w1, kernel_layout="OIW"),
|
|
relax.TensorType((2, 4, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(
|
|
x5, w4, data_layout="NCW16c", kernel_layout="IOW16i", out_layout="NWC16c"
|
|
),
|
|
relax.TensorType((2, 30, 3, 16), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x2, w0), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x3, w0), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x0, w2), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x0, w3), relax.TensorType(dtype="float32", ndim=3)
|
|
)
|
|
_check_inference(bb, relax.op.nn.conv1d_transpose(x4, w0), relax.TensorType(dtype="", ndim=3))
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
c16 = tirx.Var("c16", "int64")
|
|
iw = tirx.Var("iw", "int64")
|
|
ki = tirx.Var("ki", "int64")
|
|
ko = tirx.Var("ko", "int64")
|
|
kw = tirx.Var("kw", "int64")
|
|
x0 = relax.Var("x", R.Tensor((n, c, iw), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, c, iw, c16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((ki, ko, kw), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((c, ko, kw), "float32"))
|
|
w2 = relax.Var("w", R.Tensor((c, ko, kw, c16), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0),
|
|
relax.TensorType((n, ko, iw + kw - 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w1),
|
|
relax.TensorType((n, ko, iw + kw - 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(
|
|
x1, w2, data_layout="NCW16c", kernel_layout="IOW16i", out_layout="NCW"
|
|
),
|
|
relax.TensorType((n, ko, iw + kw - 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, strides=2, padding=1, dilation=2, output_padding=1),
|
|
relax.TensorType(
|
|
(n, ko, iw * 2 + kw * 2 - 4),
|
|
"float32",
|
|
),
|
|
)
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s2 = relax.Var("s", relax.ShapeType(ndim=3))
|
|
s3 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
w = relax.Var("w", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(bb, relax.op.nn.conv1d(x0, w), relax.TensorType(dtype="float32", ndim=3))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x1, w, data_layout="NCW16c"),
|
|
relax.TensorType(dtype="float32", ndim=4),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x0, w, out_layout="NCW16c"),
|
|
relax.TensorType(dtype="float32", ndim=4),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d(x2, w),
|
|
relax.TensorType(dtype="float32", ndim=3),
|
|
)
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_groups():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 8, 28, 16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((128, 6, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((16, 6, 3, 8), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, groups=8),
|
|
relax.TensorType((2, 48, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w1, kernel_layout="IOW8i", groups=8),
|
|
relax.TensorType((2, 48, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x1, w0, data_layout="NCW16c", groups=8),
|
|
relax.TensorType((2, 3, 30, 16), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_symbolic_groups():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x = relax.Var("x", R.Tensor((n, ic * 4, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((ic, oc, 3), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x, w0, groups=4),
|
|
relax.TensorType((n, oc * 4, 30), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_input_channel_group_incompatible():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((128, 20, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, ic, 28), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((ic - 1, oc, 3), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x0, w0, groups=6))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x1, w1, groups=6))
|
|
|
|
|
|
def test_conv1d_transpose_non_positive_group():
|
|
x = relax.Var("x", R.Tensor((2, 128, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((128, 16, 3), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d_transpose(x, w, groups=0)
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d_transpose(x, w, groups=-2)
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28), "float64"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 3), "float64"))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 28), "int8"))
|
|
w2 = relax.Var("w", R.Tensor((3, 4, 3), "int8"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 28), "int32"))
|
|
w3 = relax.Var("w", R.Tensor((3, 4, 3), "int32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x0, w0), relax.TensorType((2, 4, 30), "float16")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x1, w1), relax.TensorType((2, 4, 30), "float64")
|
|
)
|
|
_check_inference(bb, relax.op.nn.conv1d_transpose(x2, w2), relax.TensorType((2, 4, 30), "int8"))
|
|
_check_inference(
|
|
bb, relax.op.nn.conv1d_transpose(x3, w3), relax.TensorType((2, 4, 30), "int32")
|
|
)
|
|
|
|
|
|
def test_conv1d_transpose_unequal_input_channel():
|
|
bb = relax.BlockBuilder()
|
|
ic = tirx.Var("ic", "int64")
|
|
x0 = relax.Var("x", R.Tensor([2, 3, 28], "float32"))
|
|
w0 = relax.Var("w", R.Tensor([4, 3, 3], "float32"))
|
|
x1 = relax.Var("x", R.Tensor([2, ic, 28], "float32"))
|
|
w1 = relax.Var("w", R.Tensor([ic + 2, 4, 3], "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x0, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x1, w1))
|
|
|
|
|
|
def test_conv1d_transpose_wrong_output_padding():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor([2, 3, 28], "float32"))
|
|
w0 = relax.Var("w", R.Tensor([3, 4, 3], "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x0, w0, strides=2, output_padding=2))
|
|
|
|
|
|
def test_conv1d_transpose_stride_padding_dilation_int64():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3), "float32"))
|
|
conv1d = relax.op.nn.conv1d_transpose(x, w, strides=1, padding=1, dilation=1)
|
|
|
|
assert isinstance(conv1d.attrs.strides[0], int)
|
|
assert isinstance(conv1d.attrs.padding[0], int)
|
|
assert isinstance(conv1d.attrs.dilation[0], int)
|
|
|
|
|
|
def test_conv1d_transpose_wrong_strides_padding_dilation_length():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3), "float32"))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d_transpose(x, w, strides=(1, 2))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d_transpose(x, w, padding=(1, 2, 3))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv1d_transpose(x, w, dilation=(1, 2))
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_wrong_layout_string():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3), "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x, w, data_layout="IOW"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x, w, kernel_layout="NWC"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x, w, out_layout="OWI"))
|
|
|
|
|
|
def test_conv1d_transpose_dtype_mismatch():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3), "int8"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x, w))
|
|
|
|
|
|
def test_conv1d_transpose_wrong_input_ndim():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 6, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=5))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x0, w1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x0, w1, data_layout="NCW16c"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x0, w2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x1, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x2, w0))
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float32"))
|
|
x1 = relax.Var("x", relax.ShapeType((2, 3, 28)))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3), "float32"))
|
|
w1 = relax.Var("w", relax.FuncType([], R.Tensor((3, 4, 3), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x0, w1))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv1d_transpose(x1, w0))
|
|
|
|
|
|
def test_conv1d_transpose_infer_ty_mixed_precision():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28), "int8"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 3), "int8"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x0, w0, out_dtype="float32"),
|
|
relax.TensorType((2, 4, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv1d_transpose(x1, w1, out_dtype="int32"),
|
|
relax.TensorType((2, 4, 30), "int32"),
|
|
)
|
|
|
|
|
|
def test_conv2d_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 28, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
x3 = relax.Var("x", R.Tensor("float32"))
|
|
x4 = relax.Var("x", R.Tensor())
|
|
x5 = relax.Var("x", R.Tensor((2, 4, 28, 28, 16), "float32"))
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32", vdev0))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=4))
|
|
w3 = relax.Var("w", R.Tensor("float32"))
|
|
w4 = relax.Var("w", R.Tensor((48, 4, 3, 3, 16), "float32"))
|
|
w5 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32", vdev0))
|
|
|
|
_check_inference(bb, relax.op.nn.conv2d(x0, w0), relax.TensorType((2, 4, 26, 26), "float32"))
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d(x6, w5), relax.TensorType((2, 4, 26, 26), "float32", vdev0)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, out_dtype="float16"),
|
|
relax.TensorType((2, 4, 26, 26), "float16"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d(x0, w0, padding=1), relax.TensorType((2, 4, 28, 28), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, padding=[1, 2]),
|
|
relax.TensorType((2, 4, 28, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, padding=[1, 2, 3, 4]),
|
|
relax.TensorType((2, 4, 30, 32), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, strides=2),
|
|
relax.TensorType((2, 4, 13, 13), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, strides=(2, 3)),
|
|
relax.TensorType((2, 4, 13, 9), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, dilation=2),
|
|
relax.TensorType((2, 4, 24, 24), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, dilation=(2, 1)),
|
|
relax.TensorType((2, 4, 24, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x1, w0, data_layout="NHWC"),
|
|
relax.TensorType((2, 26, 26, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, out_layout="NHWC"),
|
|
relax.TensorType((2, 26, 26, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w1, kernel_layout="IOHW"),
|
|
relax.TensorType((2, 4, 26, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(
|
|
x5, w4, data_layout="NCHW16c", kernel_layout="OIHW16i", out_layout="NHWC16c"
|
|
),
|
|
relax.TensorType((2, 26, 26, 3, 16), "float32"),
|
|
)
|
|
_check_inference(bb, relax.op.nn.conv2d(x2, w0), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(bb, relax.op.nn.conv2d(x3, w0), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(bb, relax.op.nn.conv2d(x0, w2), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(bb, relax.op.nn.conv2d(x0, w3), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(bb, relax.op.nn.conv2d(x4, w0), relax.TensorType(dtype="", ndim=4))
|
|
|
|
|
|
def test_conv2d_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
c16 = tirx.Var("c16", "int64")
|
|
ih = tirx.Var("ih", "int64")
|
|
iw = tirx.Var("iw", "int64")
|
|
ki = tirx.Var("ki", "int64")
|
|
ko = tirx.Var("ko", "int64")
|
|
kh = tirx.Var("kh", "int64")
|
|
kw = tirx.Var("kw", "int64")
|
|
x0 = relax.Var("x", R.Tensor((n, c, ih, iw), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, c, ih, iw, c16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((ko, ki, kh, kw), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((ko, c, kh, kw), "float32"))
|
|
w2 = relax.Var("w", R.Tensor((ko, c, kh, kw, c16), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0),
|
|
relax.TensorType((n, ko, ih + 1 - kh, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w1),
|
|
relax.TensorType((n, ko, ih + 1 - kh, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(
|
|
x1, w2, data_layout="NCHW16c", kernel_layout="OIHW16i", out_layout="NCHW"
|
|
),
|
|
relax.TensorType((n, ko, ih + 1 - kh, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, strides=(2, 2), padding=(1, 1), dilation=(2, 2)),
|
|
relax.TensorType(
|
|
(n, ko, tvm.tirx.floordiv(ih + 3, 2) + 1 - kh, tvm.tirx.floordiv(iw + 3, 2) + 1 - kw),
|
|
"float32",
|
|
),
|
|
)
|
|
|
|
|
|
def test_conv2d_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=5))
|
|
s2 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s3 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
w = relax.Var("w", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(bb, relax.op.nn.conv2d(x0, w), relax.TensorType(dtype="float32", ndim=4))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x1, w, data_layout="NCHW16c"),
|
|
relax.TensorType(dtype="float32", ndim=5),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w, out_layout="NCHW16c"),
|
|
relax.TensorType(dtype="float32", ndim=5),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x2, w),
|
|
relax.TensorType(dtype="float32", ndim=4),
|
|
)
|
|
|
|
|
|
def test_conv2d_infer_ty_groups():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 8, 28, 28, 16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((48, 16, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((48, 2, 3, 3, 8), "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d(x0, w0, groups=8), relax.TensorType((2, 48, 26, 26), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w1, kernel_layout="OIHW8i", groups=8),
|
|
relax.TensorType((2, 48, 26, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x1, w0, data_layout="NCHW16c", groups=8),
|
|
relax.TensorType((2, 3, 26, 26, 16), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv2d_infer_ty_symbolic_groups():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x = relax.Var("x", R.Tensor((n, ic * 4, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((oc * 4, ic, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((oc, ic, 3, 3), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x, w0, groups=4),
|
|
relax.TensorType((n, oc * 4, 26, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d(x, w1, groups=4), relax.TensorType((n, oc, 26, 26), "float32")
|
|
)
|
|
|
|
|
|
def test_conv2d_infer_ty_input_channel_group_incompatible():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((48, 20, 3, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, ic * 6, 28, 28), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((oc, ic - 1, 3, 3), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x0, w0, groups=6))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x1, w1, groups=6))
|
|
|
|
|
|
def test_conv2d_infer_ty_output_channel_group_incompatible():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x0 = relax.Var("x", R.Tensor((2, 120, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((128, 20, 3, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, ic * 6, 28, 28), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((oc * 6 + 4, ic * 6, 3, 3), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x0, w0, groups=6))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x1, w1, groups=6))
|
|
|
|
|
|
def test_conv2d_non_positive_group():
|
|
x = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((48, 16, 3, 3), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d(x, w, groups=0)
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d(x, w, groups=-2)
|
|
|
|
|
|
def test_conv2d_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float64"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float64"))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
|
|
w2 = relax.Var("w", R.Tensor((4, 3, 3, 3), "int8"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int32"))
|
|
w3 = relax.Var("w", R.Tensor((4, 3, 3, 3), "int32"))
|
|
|
|
_check_inference(bb, relax.op.nn.conv2d(x0, w0), relax.TensorType((2, 4, 26, 26), "float16"))
|
|
_check_inference(bb, relax.op.nn.conv2d(x1, w1), relax.TensorType((2, 4, 26, 26), "float64"))
|
|
_check_inference(bb, relax.op.nn.conv2d(x2, w2), relax.TensorType((2, 4, 26, 26), "int8"))
|
|
_check_inference(bb, relax.op.nn.conv2d(x3, w3), relax.TensorType((2, 4, 26, 26), "int32"))
|
|
|
|
|
|
def test_conv2d_infer_ty_mixed_precision():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 3, 3), "int8"))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 28, 28)))
|
|
w2 = relax.Var("w", R.Tensor((4, 3, 3, 3)))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x0, w0, out_dtype="float32"),
|
|
relax.TensorType((2, 4, 26, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x1, w1, out_dtype="int32"),
|
|
relax.TensorType((2, 4, 26, 26), "int32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d(x2, w2, out_dtype="float32"),
|
|
relax.TensorType((2, 4, 26, 26), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv2d_unequal_input_channel():
|
|
bb = relax.BlockBuilder()
|
|
ic = tirx.Var("ic", "int64")
|
|
x0 = relax.Var("x", R.Tensor([2, 3, 28, 28], "float32"))
|
|
w0 = relax.Var("w", R.Tensor([3, 4, 3, 3], "float32"))
|
|
x1 = relax.Var("x", R.Tensor([2, ic, 28, 28], "float32"))
|
|
w1 = relax.Var("w", R.Tensor([4, ic + 2, 3, 3], "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x0, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x1, w1))
|
|
|
|
|
|
def test_conv2d_stride_padding_dilation_int64():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
conv2d = relax.op.nn.conv2d(x, w, strides=(1, 1), padding=(1, 1), dilation=(1, 1))
|
|
|
|
assert isinstance(conv2d.attrs.strides[0], int)
|
|
assert isinstance(conv2d.attrs.strides[1], int)
|
|
assert isinstance(conv2d.attrs.padding[0], int)
|
|
assert isinstance(conv2d.attrs.padding[1], int)
|
|
assert isinstance(conv2d.attrs.padding[2], int)
|
|
assert isinstance(conv2d.attrs.padding[3], int)
|
|
assert isinstance(conv2d.attrs.dilation[0], int)
|
|
assert isinstance(conv2d.attrs.dilation[1], int)
|
|
|
|
|
|
def test_conv2d_wrong_strides_padding_dilation_length():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d(x, w, strides=(1, 2, 3))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d(x, w, padding=(1, 2, 3))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d(x, w, dilation=(1, 2, 3))
|
|
|
|
|
|
def test_conv2d_infer_ty_wrong_layout_string():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x, w, data_layout="OIHW"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x, w, kernel_layout="NHWC"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x, w, out_layout="OHWI"))
|
|
|
|
|
|
def test_conv2d_dtype_mismatch():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((4, 3, 3, 3), "int8"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv2d(x, w))
|
|
|
|
|
|
def test_conv2d_wrong_input_ndim():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 6, 3, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=6))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x0, w1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x0, w1, data_layout="NCHW16c"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x0, w2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x1, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d(x2, w0))
|
|
|
|
|
|
def test_conv2d_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
x1 = relax.Var("x", relax.ShapeType((2, 3, 28, 28)))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
w1 = relax.Var("w", relax.FuncType([], R.Tensor((4, 3, 3, 3), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv2d(x0, w1))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv2d(x1, w0))
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 28, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=4))
|
|
x3 = relax.Var("x", R.Tensor("float32"))
|
|
x4 = relax.Var("x", R.Tensor())
|
|
x5 = relax.Var("x", R.Tensor((2, 4, 28, 28, 16), "float32"))
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32", vdev0))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((4, 3, 3, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=4))
|
|
w3 = relax.Var("w", R.Tensor("float32"))
|
|
w4 = relax.Var("w", R.Tensor((4, 48, 3, 3, 16), "float32"))
|
|
w5 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x0, w0), relax.TensorType((2, 4, 30, 30), "float32")
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x6, w5),
|
|
relax.TensorType((2, 4, 30, 30), "float32", vdev0),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, out_dtype="float16"),
|
|
relax.TensorType((2, 4, 30, 30), "float16"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, padding=1),
|
|
relax.TensorType((2, 4, 28, 28), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, padding=[1, 2]),
|
|
relax.TensorType((2, 4, 28, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, padding=[1, 2, 3, 4]),
|
|
relax.TensorType((2, 4, 26, 24), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, strides=3, output_padding=1),
|
|
relax.TensorType((2, 4, 85, 85), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, strides=3, output_padding=[2, 1]),
|
|
relax.TensorType((2, 4, 86, 85), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, strides=2),
|
|
relax.TensorType((2, 4, 57, 57), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, strides=(2, 3)),
|
|
relax.TensorType((2, 4, 57, 84), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, dilation=2),
|
|
relax.TensorType((2, 4, 32, 32), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, dilation=(2, 1)),
|
|
relax.TensorType((2, 4, 32, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x1, w0, data_layout="NHWC"),
|
|
relax.TensorType((2, 30, 30, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, out_layout="NHWC"),
|
|
relax.TensorType((2, 30, 30, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w1, kernel_layout="OIHW"),
|
|
relax.TensorType((2, 4, 30, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(
|
|
x5, w4, data_layout="NCHW16c", kernel_layout="IOHW16i", out_layout="NHWC16c"
|
|
),
|
|
relax.TensorType((2, 30, 30, 3, 16), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x2, w0), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x3, w0), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x0, w2), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x0, w3), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(bb, relax.op.nn.conv2d_transpose(x4, w0), relax.TensorType(dtype="", ndim=4))
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
c16 = tirx.Var("c16", "int64")
|
|
ih = tirx.Var("ih", "int64")
|
|
iw = tirx.Var("iw", "int64")
|
|
ki = tirx.Var("ki", "int64")
|
|
ko = tirx.Var("ko", "int64")
|
|
kh = tirx.Var("kh", "int64")
|
|
kw = tirx.Var("kw", "int64")
|
|
x0 = relax.Var("x", R.Tensor((n, c, ih, iw), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, c, ih, iw, c16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((ki, ko, kh, kw), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((c, ko, kh, kw), "float32"))
|
|
w2 = relax.Var("w", R.Tensor((c, ko, kh, kw, c16), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0),
|
|
relax.TensorType((n, ko, ih + kh - 1, iw + kw - 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w1),
|
|
relax.TensorType((n, ko, ih + kh - 1, iw + kw - 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(
|
|
x1, w2, data_layout="NCHW16c", kernel_layout="IOHW16i", out_layout="NCHW"
|
|
),
|
|
relax.TensorType((n, ko, ih + kh - 1, iw + kw - 1), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(
|
|
x0, w0, strides=(2, 2), padding=(1, 1), output_padding=(1, 0), dilation=(2, 2)
|
|
),
|
|
relax.TensorType(
|
|
(n, ko, ih * 2 + kh * 2 - 4, iw * 2 + kw * 2 - 5),
|
|
"float32",
|
|
),
|
|
)
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=5))
|
|
s2 = relax.Var("s", relax.ShapeType(ndim=4))
|
|
s3 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
w = relax.Var("w", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x0, w), relax.TensorType(dtype="float32", ndim=4)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x1, w, data_layout="NCHW16c"),
|
|
relax.TensorType(dtype="float32", ndim=5),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w, out_layout="NCHW16c"),
|
|
relax.TensorType(dtype="float32", ndim=5),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x2, w),
|
|
relax.TensorType(dtype="float32", ndim=4),
|
|
)
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_groups():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 8, 28, 28, 16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((128, 6, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((16, 6, 3, 3, 8), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, groups=8),
|
|
relax.TensorType((2, 48, 30, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w1, kernel_layout="IOHW8i", groups=8),
|
|
relax.TensorType((2, 48, 30, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x1, w0, data_layout="NCHW16c", groups=8),
|
|
relax.TensorType((2, 3, 30, 30, 16), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_symbolic_groups():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x = relax.Var("x", R.Tensor((n, ic * 4, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((ic, oc, 3, 3), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x, w0, groups=4),
|
|
relax.TensorType((n, oc * 4, 30, 30), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_input_channel_group_incompatible():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
ic = tirx.Var("c", "int64")
|
|
oc = tirx.Var("oc", "int64")
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((128, 20, 3, 3), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, ic, 28, 28), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((ic - 1, oc, 3, 3), "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w0, groups=6))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x1, w1, groups=6))
|
|
|
|
|
|
def test_conv2d_transpose_non_positive_group():
|
|
x = relax.Var("x", R.Tensor((2, 128, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((128, 16, 3, 3), "float32"))
|
|
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d_transpose(x, w, groups=0)
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d_transpose(x, w, groups=-2)
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_more_input_dtype():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float64"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float64"))
|
|
x2 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
|
|
w2 = relax.Var("w", R.Tensor((3, 4, 3, 3), "int8"))
|
|
x3 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int32"))
|
|
w3 = relax.Var("w", R.Tensor((3, 4, 3, 3), "int32"))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x0, w0), relax.TensorType((2, 4, 30, 30), "float16")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x1, w1), relax.TensorType((2, 4, 30, 30), "float64")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x2, w2), relax.TensorType((2, 4, 30, 30), "int8")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv2d_transpose(x3, w3), relax.TensorType((2, 4, 30, 30), "int32")
|
|
)
|
|
|
|
|
|
def test_conv2d_transpose_unequal_input_channel():
|
|
bb = relax.BlockBuilder()
|
|
ic = tirx.Var("ic", "int64")
|
|
x0 = relax.Var("x", R.Tensor([2, 3, 28, 28], "float32"))
|
|
w0 = relax.Var("w", R.Tensor([4, 3, 3, 3], "float32"))
|
|
x1 = relax.Var("x", R.Tensor([2, ic, 28, 28], "float32"))
|
|
w1 = relax.Var("w", R.Tensor([ic + 2, 4, 3, 3], "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x1, w1))
|
|
|
|
|
|
def test_conv2d_transpose_wrong_output_padding():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor([2, 3, 28, 28], "float32"))
|
|
w0 = relax.Var("w", R.Tensor([3, 4, 3, 3], "float32"))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w0, strides=2, output_padding=2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w0, strides=(2, 2), output_padding=(2, 2)))
|
|
|
|
|
|
def test_conv2d_transpose_stride_padding_dilation_int64():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
|
|
conv2d_transpose = relax.op.nn.conv2d_transpose(
|
|
x, w, strides=(1, 1), padding=(1, 1), output_padding=(1, 2), dilation=(1, 1)
|
|
)
|
|
|
|
assert isinstance(conv2d_transpose.attrs.strides[0], int)
|
|
assert isinstance(conv2d_transpose.attrs.strides[1], int)
|
|
assert isinstance(conv2d_transpose.attrs.padding[0], int)
|
|
assert isinstance(conv2d_transpose.attrs.padding[1], int)
|
|
assert isinstance(conv2d_transpose.attrs.padding[2], int)
|
|
assert isinstance(conv2d_transpose.attrs.padding[3], int)
|
|
assert isinstance(conv2d_transpose.attrs.output_padding[0], int)
|
|
assert isinstance(conv2d_transpose.attrs.output_padding[1], int)
|
|
assert isinstance(conv2d_transpose.attrs.dilation[0], int)
|
|
assert isinstance(conv2d_transpose.attrs.dilation[1], int)
|
|
|
|
|
|
def test_conv2d_transpose_wrong_strides_padding_dilation_length():
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d_transpose(x, w, strides=(1, 2, 3))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d_transpose(x, w, padding=(1, 2, 3))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d_transpose(x, w, output_padding=(1, 2, 3))
|
|
with pytest.raises(tvm.error.InternalError):
|
|
relax.op.nn.conv2d_transpose(x, w, dilation=(1, 2, 3))
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_wrong_layout_string():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x, w, data_layout="IOHW"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x, w, kernel_layout="NHWC"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x, w, out_layout="OHWI"))
|
|
|
|
|
|
def test_conv2d_transpose_dtype_mismatch():
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
w = relax.Var("w", R.Tensor((3, 4, 3, 3), "int8"))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x, w))
|
|
|
|
|
|
def test_conv2d_transpose_wrong_input_ndim():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=3))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 6, 3, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=6))
|
|
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w1))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w1, data_layout="NCHW16c"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x1, w0))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x2, w0))
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_wrong_input_type():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float32"))
|
|
x1 = relax.Var("x", relax.ShapeType((2, 3, 28, 28)))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float32"))
|
|
w1 = relax.Var("w", relax.FuncType([], R.Tensor((3, 4, 3, 3), "float32")))
|
|
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x0, w1))
|
|
with pytest.raises(TypeError):
|
|
bb.normalize(relax.op.nn.conv2d_transpose(x1, w0))
|
|
|
|
|
|
def test_conv2d_transpose_infer_ty_mixed_precision():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28), "float16"))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3), "float16"))
|
|
x1 = relax.Var("x", R.Tensor((2, 3, 28, 28), "int8"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 3, 3), "int8"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x0, w0, out_dtype="float32"),
|
|
relax.TensorType((2, 4, 30, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv2d_transpose(x1, w1, out_dtype="int32"),
|
|
relax.TensorType((2, 4, 30, 30), "int32"),
|
|
)
|
|
|
|
|
|
def test_conv3d_transpose_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3, 3), "float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d_transpose(x0, w0),
|
|
relax.TensorType((2, 4, 30, 30, 30), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d_transpose(x0, w0, padding=1),
|
|
relax.TensorType((2, 4, 28, 28, 28), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d_transpose(x0, w0, strides=2, output_padding=1),
|
|
relax.TensorType((2, 4, 58, 58, 58), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv3d_transpose_infer_ty_ndhwc_out_layout():
|
|
bb = relax.BlockBuilder()
|
|
x_ndhwc = relax.Var("x_nd", R.Tensor((2, 28, 28, 28, 3), "float32"))
|
|
x_ncdhw = relax.Var("x_nc", R.Tensor((2, 3, 28, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3, 3), "float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d_transpose(x_ndhwc, w0, data_layout="NDHWC"),
|
|
relax.TensorType((2, 30, 30, 30, 4), "float32"),
|
|
)
|
|
# Default data_layout is NCDHW; use NCDHW-shaped input when only out_layout is NDHWC.
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d_transpose(x_ncdhw, w0, out_layout="NDHWC"),
|
|
relax.TensorType((2, 30, 30, 30, 4), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv3d_transpose_infer_ty_groups():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 128, 28, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((128, 16, 3, 3, 3), "float32"))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d_transpose(x0, w0, groups=8),
|
|
relax.TensorType((2, 128, 30, 30, 30), "float32"),
|
|
)
|
|
|
|
|
|
def test_conv3d_transpose_wrong_output_padding():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((3, 4, 3, 3, 3), "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv3d_transpose(x0, w0, strides=2, output_padding=2))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(
|
|
relax.op.nn.conv3d_transpose(x0, w0, strides=(2, 2, 2), output_padding=(2, 2, 2))
|
|
)
|
|
|
|
|
|
def test_conv3d_transpose_unequal_input_channel():
|
|
bb = relax.BlockBuilder()
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28, 28), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((4, 4, 3, 3, 3), "float32"))
|
|
with pytest.raises(ValueError):
|
|
bb.normalize(relax.op.nn.conv3d_transpose(x0, w0))
|
|
|
|
|
|
def test_conv3d_infer_ty():
|
|
bb = relax.BlockBuilder()
|
|
vdev0 = VDevice("llvm")
|
|
x0 = relax.Var("x", R.Tensor((2, 3, 28, 28, 28), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((2, 28, 28, 28, 3), "float32"))
|
|
x2 = relax.Var("x", R.Tensor("float32", ndim=5))
|
|
x3 = relax.Var("x", R.Tensor("float32"))
|
|
x4 = relax.Var("x", R.Tensor())
|
|
x5 = relax.Var("x", R.Tensor((2, 4, 28, 28, 28, 16), "float32"))
|
|
x6 = relax.Var("x", R.Tensor((2, 3, 28, 28, 28), "float32", vdev0))
|
|
w0 = relax.Var("w", R.Tensor((4, 3, 3, 3, 3), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((3, 4, 3, 3, 3), "float32"))
|
|
w2 = relax.Var("w", R.Tensor("float32", ndim=5))
|
|
w3 = relax.Var("w", R.Tensor("float32"))
|
|
w4 = relax.Var("w", R.Tensor((48, 4, 3, 3, 3, 16), "float32"))
|
|
w5 = relax.Var("w", R.Tensor((4, 3, 3, 3, 3), "float32", vdev0))
|
|
|
|
_check_inference(
|
|
bb, relax.op.nn.conv3d(x0, w0), relax.TensorType((2, 4, 26, 26, 26), "float32")
|
|
)
|
|
_check_inference(
|
|
bb, relax.op.nn.conv3d(x6, w5), relax.TensorType((2, 4, 26, 26, 26), "float32", vdev0)
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, out_dtype="float16"),
|
|
relax.TensorType((2, 4, 26, 26, 26), "float16"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, padding=1),
|
|
relax.TensorType((2, 4, 28, 28, 28), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, padding=[1, 2, 3]),
|
|
relax.TensorType((2, 4, 28, 30, 32), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, padding=[1, 2, 3, 4, 5, 6]),
|
|
relax.TensorType((2, 4, 31, 33, 35), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, strides=2),
|
|
relax.TensorType((2, 4, 13, 13, 13), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, strides=(2, 3, 4)),
|
|
relax.TensorType((2, 4, 13, 9, 7), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, dilation=2),
|
|
relax.TensorType((2, 4, 24, 24, 24), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, dilation=(3, 2, 1)),
|
|
relax.TensorType((2, 4, 22, 24, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x1, w0, data_layout="NDHWC"),
|
|
relax.TensorType((2, 26, 26, 26, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, out_layout="NDHWC"),
|
|
relax.TensorType((2, 26, 26, 26, 4), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w1, kernel_layout="IODHW"),
|
|
relax.TensorType((2, 4, 26, 26, 26), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(
|
|
x5, w4, data_layout="NCDHW16c", kernel_layout="OIDHW16i", out_layout="NDHWC16c"
|
|
),
|
|
relax.TensorType((2, 26, 26, 26, 3, 16), "float32"),
|
|
)
|
|
_check_inference(bb, relax.op.nn.conv3d(x2, w0), relax.TensorType(dtype="float32", ndim=5))
|
|
_check_inference(bb, relax.op.nn.conv3d(x3, w0), relax.TensorType(dtype="float32", ndim=5))
|
|
_check_inference(bb, relax.op.nn.conv3d(x0, w2), relax.TensorType(dtype="float32", ndim=5))
|
|
_check_inference(bb, relax.op.nn.conv3d(x0, w3), relax.TensorType(dtype="float32", ndim=5))
|
|
_check_inference(bb, relax.op.nn.conv3d(x4, w0), relax.TensorType(dtype="", ndim=5))
|
|
|
|
|
|
def test_conv3d_infer_ty_shape_symbolic():
|
|
bb = relax.BlockBuilder()
|
|
n = tirx.Var("n", "int64")
|
|
c = tirx.Var("c", "int64")
|
|
c16 = tirx.Var("c16", "int64")
|
|
id = tirx.Var("id", "int64")
|
|
ih = tirx.Var("ih", "int64")
|
|
iw = tirx.Var("iw", "int64")
|
|
ki = tirx.Var("ki", "int64")
|
|
ko = tirx.Var("ko", "int64")
|
|
kd = tirx.Var("kd", "int64")
|
|
kh = tirx.Var("kh", "int64")
|
|
kw = tirx.Var("kw", "int64")
|
|
x0 = relax.Var("x", R.Tensor((n, c, id, ih, iw), "float32"))
|
|
x1 = relax.Var("x", R.Tensor((n, c, id, ih, iw, c16), "float32"))
|
|
w0 = relax.Var("w", R.Tensor((ko, ki, kd, kh, kw), "float32"))
|
|
w1 = relax.Var("w", R.Tensor((ko, c, kd, kh, kw), "float32"))
|
|
w2 = relax.Var("w", R.Tensor((ko, c, kd, kh, kw, c16), "float32"))
|
|
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0),
|
|
relax.TensorType((n, ko, id + 1 - kd, ih + 1 - kh, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w1),
|
|
relax.TensorType((n, ko, id + 1 - kd, ih + 1 - kh, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(
|
|
x1, w2, data_layout="NCDHW16c", kernel_layout="OIDHW16i", out_layout="NCDHW"
|
|
),
|
|
relax.TensorType((n, ko, id + 1 - kd, ih + 1 - kh, iw + 1 - kw), "float32"),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w0, strides=(2, 2, 2), padding=(1, 1, 1), dilation=(2, 2, 2)),
|
|
relax.TensorType(
|
|
(
|
|
n,
|
|
ko,
|
|
tvm.tirx.floordiv(id + 3, 2) + 1 - kd,
|
|
tvm.tirx.floordiv(ih + 3, 2) + 1 - kh,
|
|
tvm.tirx.floordiv(iw + 3, 2) + 1 - kw,
|
|
),
|
|
"float32",
|
|
),
|
|
)
|
|
|
|
|
|
def test_conv3d_infer_ty_shape_var():
|
|
bb = relax.BlockBuilder()
|
|
s0 = relax.Var("s", relax.ShapeType(ndim=5))
|
|
s1 = relax.Var("s", relax.ShapeType(ndim=6))
|
|
s2 = relax.Var("s", relax.ShapeType(ndim=5))
|
|
s3 = relax.Var("s", relax.ShapeType())
|
|
x0 = relax.Var("x", relax.TensorType(s0, "float32"))
|
|
x1 = relax.Var("x", relax.TensorType(s1, "float32"))
|
|
x2 = relax.Var("x", relax.TensorType(s3, "float32"))
|
|
w = relax.Var("w", relax.TensorType(s2, "float32"))
|
|
|
|
_check_inference(bb, relax.op.nn.conv3d(x0, w), relax.TensorType(dtype="float32", ndim=5))
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x1, w, data_layout="NCDHW16c"),
|
|
relax.TensorType(dtype="float32", ndim=6),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x0, w, out_layout="NCDHW16c"),
|
|
relax.TensorType(dtype="float32", ndim=6),
|
|
)
|
|
_check_inference(
|
|
bb,
|
|
relax.op.nn.conv3d(x2, w),
|
|
relax.TensorType(dtype="float32", ndim=5),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|