Skip to content

Commit

Permalink
add back missed conv2d_output_shape by mistake
Browse files Browse the repository at this point in the history
  • Loading branch information
optima2005 committed Nov 22, 2019
1 parent c82e320 commit 536137a
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 0 deletions.
65 changes: 65 additions & 0 deletions python/tvm/contrib/cudnn.py
Expand Up @@ -172,6 +172,70 @@ def conv2d_w_shape(in_channel,
"""
return [out_channel, in_channel, filter_h, filter_w]

def conv2d_output_shape(tensor_format,
pad_h,
pad_w,
stride_h,
stride_w,
dilation_h,
dilation_w,
x_shape,
w_shape):
"""Get output shape of 2D convolution
Paramters
---------
tensor_format: int
0: CUDNN_TENSOR_NCHW
1: CUDNN_TENSOR_NHWC
2: CUDNN_TENSOR_NCHW_VECT_C
pad_h: int
height pad
pad_w: int
weight pad
stride_h: int
height stride
stride_w: int
width stride
dilation_h: int
height dilation
dilation_w: int
width dilation
x_shape: list
input shape
w_shape: list
weight shape
Returns
-------
oshape: list
output shape
"""
assert isinstance(x_shape, list)
assert isinstance(w_shape, list)
assert len(x_shape) == 4
assert len(w_shape) == 4
oshape = np.zeros((len(x_shape)), dtype=np.int32)
func = _get_global_func("tvm.contrib.cudnn.conv2d.output_shape")
func(tensor_format,
pad_h,
pad_w,
stride_h,
stride_w,
dilation_h,
dilation_w,
x_shape[0].value,
x_shape[1].value,
x_shape[2].value,
x_shape[3].value,
w_shape[0].value,
w_shape[1].value,
w_shape[2].value,
w_shape[3].value,
_get_np_int32_array_handle(oshape))
return list(oshape)


def conv2d_find_algo(tensor_format,
pad_h,
pad_w,
Expand Down Expand Up @@ -321,6 +385,7 @@ def conv2d_forward(x,
ins[1],
outs[0]), name="y")


def conv3d_w_shape(in_channel,
out_channel,
filter_d,
Expand Down
1 change: 1 addition & 0 deletions tests/python/relay/test_op_level2.py
Expand Up @@ -867,6 +867,7 @@ def test_bitpack_infer_type():
test_conv2d_transpose_run()
test_conv2d_run()
test_conv2d_winograd()
test_conv3d_run()
test_bitserial_conv2d_infer_type()
test_batch_flatten()
test_upsampling()
Expand Down

0 comments on commit 536137a

Please sign in to comment.