This is an automated email from the ASF dual-hosted git repository.
jcf94 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new f64ddff Add missing shape functions for relay.nn operations (#8489)
f64ddff is described below
commit f64ddff4977cbb3edc8391a5ee4802e4870846c2
Author: Rijul Gupta <[email protected]>
AuthorDate: Mon Jul 19 12:39:29 2021 +1000
Add missing shape functions for relay.nn operations (#8489)
* Update _nn.py
add a few missing shape functions
* Update _nn.py
Updated conv_transpose shape function to accomodate conv1d_transpose
* added tests for new functions
* fixed a lint error
* fixed shape func error
* attempt fixing cuda error
---
python/tvm/relay/op/nn/_nn.py | 31 +++----
python/tvm/topi/cuda/conv1d_transpose_ncw.py | 2 +-
tests/python/relay/test_any.py | 123 +++++++++++++++++++++++++++
3 files changed, 140 insertions(+), 16 deletions(-)
diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py
index b0b5a56..6e757ea 100644
--- a/python/tvm/relay/op/nn/_nn.py
+++ b/python/tvm/relay/op/nn/_nn.py
@@ -1082,27 +1082,22 @@ def conv2d_NCHWc_shape_func(attrs, inputs, _):
@script
-def _conv2d_transpose_nchw_shape_func(dshape, kshape, strides, padding,
dilation, output_padding):
+def _conv_transpose_shape_func(dshape, kshape, strides, padding, dilation,
output_padding):
out = output_tensor((dshape.shape[0],), "int64")
- kheight = kshape[2]
- kwidth = kshape[3]
- dilated_kh = (kheight - 1) * dilation[0] + 1
- dilated_kw = (kwidth - 1) * dilation[1] + 1
-
- out_height = strides[0] * (dshape[2] - 1) + dilated_kh - 2 * padding[0] +
output_padding[0]
- out_width = strides[1] * (dshape[3] - 1) + dilated_kw - 2 * padding[1] +
output_padding[1]
-
out[0] = dshape[0]
out[1] = kshape[1]
- out[2] = out_height
- out[3] = out_width
+
+ for i in const_range(dshape.shape[0] - 2):
+ dilated_k = (kshape[i + 2] - 1) * dilation[i] + 1
+ out[i + 2] = (
+ strides[i] * (dshape[i + 2] - 1) + dilated_k - 2 * padding[i] +
output_padding[i]
+ )
return out
[email protected]_shape_func("nn.conv2d_transpose", False)
-def conv2d_transpose_nchw_shape_func(attrs, inputs, _):
+def conv_transpose_shape_func(attrs, inputs, _):
"""
- Shape function for conv2d_transpose op.
+ Shape function for contrib_conv2d_NCHWc op.
"""
strides = get_const_tuple(attrs.strides)
padding = get_const_tuple(attrs.padding)
@@ -1110,7 +1105,7 @@ def conv2d_transpose_nchw_shape_func(attrs, inputs, _):
output_padding = get_const_tuple(attrs.output_padding)
return [
- _conv2d_transpose_nchw_shape_func(
+ _conv_transpose_shape_func(
inputs[0],
inputs[1],
convert(strides),
@@ -1121,6 +1116,10 @@ def conv2d_transpose_nchw_shape_func(attrs, inputs, _):
]
+reg.register_shape_func("nn.conv1d_transpose", False,
conv_transpose_shape_func)
+reg.register_shape_func("nn.conv2d_transpose", False,
conv_transpose_shape_func)
+
+
@script
def _pool2d_shape_func(data_shape, pool_size, strides, padding, height_axis,
width_axis):
out = output_tensor((data_shape.shape[0],), "int64")
@@ -1339,3 +1338,5 @@ reg.register_shape_func("nn.bias_add", False,
elemwise_shape_func)
reg.register_shape_func("nn.softmax", False, elemwise_shape_func)
reg.register_shape_func("nn.fast_softmax", False, elemwise_shape_func)
reg.register_shape_func("nn.relu", False, elemwise_shape_func)
+reg.register_shape_func("nn.leaky_relu", False, elemwise_shape_func)
+reg.register_shape_func("nn.prelu", False, elemwise_shape_func)
diff --git a/python/tvm/topi/cuda/conv1d_transpose_ncw.py
b/python/tvm/topi/cuda/conv1d_transpose_ncw.py
index 58f53ea..2098aa9 100644
--- a/python/tvm/topi/cuda/conv1d_transpose_ncw.py
+++ b/python/tvm/topi/cuda/conv1d_transpose_ncw.py
@@ -142,7 +142,7 @@ def schedule_conv1d_transpose_ncw(cfg, outs):
##### space definition begin #####
n, f, x = s[conv].op.axis
rc = s[conv].op.reduce_axis[0]
- cfg.define_split("tile_n", cfg.axis(n), num_outputs=4)
+ cfg.define_split("tile_n", cfg.axis(n if isinstance(n, int) else
1), num_outputs=4)
cfg.define_split("tile_f", cfg.axis(f), num_outputs=4)
cfg.define_split("tile_x", cfg.axis(x), num_outputs=4)
cfg.define_split("tile_rc", cfg.axis(rc), num_outputs=3)
diff --git a/tests/python/relay/test_any.py b/tests/python/relay/test_any.py
index f871c67..7b1d416 100644
--- a/tests/python/relay/test_any.py
+++ b/tests/python/relay/test_any.py
@@ -643,6 +643,63 @@ def test_any_conv2d_NCHWc():
)
+def verify_any_conv1d_transpose_ncw(
+ data_shape,
+ kernel_shape,
+ strides,
+ padding,
+ dilation,
+ groups,
+ static_data_shape,
+ ref_out_shape,
+ output_padding,
+):
+ mod = tvm.IRModule()
+ dtype = "float32"
+ data = relay.var("data", shape=data_shape, dtype=dtype)
+ kernel = relay.var("kernel", shape=kernel_shape, dtype=dtype)
+ y = relay.nn.conv1d_transpose(
+ data,
+ kernel,
+ strides,
+ padding,
+ dilation,
+ groups,
+ kernel_size=kernel_shape[2:],
+ output_padding=output_padding,
+ )
+ mod["main"] = relay.Function([data, kernel], y)
+ data_np = np.random.uniform(size=static_data_shape).astype(dtype)
+ kernel_np = np.random.uniform(size=kernel_shape).astype(dtype)
+ check_result([data_np, kernel_np], mod, ref_out_shape, assert_shape=True)
+
+
[email protected]_gpu
+def test_any_conv1d_transpose_ncw():
+ verify_any_conv1d_transpose_ncw(
+ (relay.Any(), 64, 224),
+ (64, 192, 3),
+ (1,),
+ (1,),
+ (1,),
+ 1,
+ (2, 64, 224),
+ (2, 192, 224),
+ (0, 0),
+ )
+ verify_any_conv1d_transpose_ncw(
+ (relay.Any(), 32, 224),
+ (32, 64, 3),
+ (2,),
+ (1,),
+ (1,),
+ 1,
+ (1, 32, 224),
+ (1, 64, 448),
+ (1, 1),
+ )
+
+
def verify_any_conv2d_transpose_nchw(
data_shape,
kernel_shape,
@@ -932,6 +989,72 @@ def test_any_softmax():
verify_any_softmax(any_dims(4), 2, (13, 11, 3, 1), (13, 11, 3, 1))
+def verify_any_relu(data_shape, static_data_shape, ref_out_shape):
+ mod = tvm.IRModule()
+ dtype = "float32"
+ data = relay.var("data", shape=data_shape, dtype=dtype)
+ y = relay.nn.relu(data)
+ mod["main"] = relay.Function([data], y)
+ data_np = np.random.uniform(size=static_data_shape).astype(dtype)
+ check_result([data_np], mod, ref_out_shape, assert_shape=True)
+
+
[email protected]_gpu
+def test_any_relu():
+ verify_any_relu(any_dims(3), (1, 2, 3), (1, 2, 3))
+ verify_any_relu(any_dims(4), (13, 11, 3, 1), (13, 11, 3, 1))
+
+
+def verify_any_prelu(data_shape, alpha, static_data_shape, ref_out_shape):
+ mod = tvm.IRModule()
+ dtype = "float32"
+ data = relay.var("data", shape=data_shape, dtype=dtype)
+ alpha = relay.const(np.array([alpha]), dtype=dtype)
+ y = relay.nn.prelu(data, alpha)
+ mod["main"] = relay.Function([data], y)
+ data_np = np.random.uniform(size=static_data_shape).astype(dtype)
+ check_result([data_np], mod, ref_out_shape, assert_shape=True)
+
+
[email protected]_gpu
+def test_any_prelu():
+ verify_any_prelu(any_dims(3), 1, (1, 2, 3), (1, 2, 3))
+ verify_any_prelu(any_dims(4), 2, (13, 11, 3, 1), (13, 11, 3, 1))
+
+
+def verify_any_leaky_relu(data_shape, alpha, static_data_shape, ref_out_shape):
+ mod = tvm.IRModule()
+ dtype = "float32"
+ data = relay.var("data", shape=data_shape, dtype=dtype)
+ y = relay.nn.leaky_relu(data, alpha)
+ mod["main"] = relay.Function([data], y)
+ data_np = np.random.uniform(size=static_data_shape).astype(dtype)
+ check_result([data_np], mod, ref_out_shape, assert_shape=True)
+
+
[email protected]_gpu
+def test_any_leaky_relu():
+ verify_any_leaky_relu(any_dims(3), 0.1, (1, 2, 3), (1, 2, 3))
+ verify_any_leaky_relu(any_dims(4), 0.2, (13, 11, 3, 1), (13, 11, 3, 1))
+
+
+def verify_any_bias_add(data_shape, static_data_shape, ref_out_shape):
+ mod = tvm.IRModule()
+ dtype = "float32"
+ data = relay.var("data", shape=data_shape, dtype=dtype)
+ bias = relay.const(np.random.randn(1), dtype=dtype)
+ y = relay.nn.bias_add(data, bias)
+ mod["main"] = relay.Function([data], y)
+ data_np = np.random.uniform(size=static_data_shape).astype(dtype)
+ check_result([data_np], mod, ref_out_shape, assert_shape=True)
+
+
[email protected]_gpu
+def test_any_bias_add():
+ verify_any_bias_add(any_dims(3), (1, 2, 3), (1, 2, 3))
+ verify_any_bias_add(any_dims(4), (13, 11, 3, 1), (13, 11, 3, 1))
+
+
def verify_any_topk(data_shape, kval, np_dshape, dtype, ret_type="indices",
const_k=False):
mod = tvm.IRModule()
data = relay.var("data", shape=data_shape, dtype=dtype)
