def verify_strided_slice(in_shape, begin, end, strides=None):
A = te.placeholder(shape=in_shape, name="A")
strides = [1, 1, 1] if strides is None else strides
B = topi.strided_slice(A, begin, end, strides) + 1
def check_device(device):
ctx = tvm.context(device, 0)
if not tvm.testing.device_enabled(device):
print("Skip because %s is not enabled" % device)
return
print("Running on target: %s" % device)
with tvm.target.Target(device):
s = tvm.topi.testing.get_injective_schedule(device)(B)
foo = tvm.build(s, [A, B], device, name="stride_slice")
x_np = np.random.uniform(size=in_shape).astype(A.dtype)
out_npy = tvm.topi.testing.strided_slice_python(
x_np, begin, end, strides) + 1
data_nd = tvm.nd.array(x_np, ctx)
out_nd = tvm.nd.empty(out_npy.shape, ctx=ctx, dtype=A.dtype)
foo(data_nd, out_nd)
tvm.testing.assert_allclose(out_nd.asnumpy(), out_npy)
for device in ["llvm", "opencl", "sdaccel", "aocl_sw_emu"]:
check_device(device)
def conv_bwd(N, CI, HI, WI, CO, HO, WO, KSIZE, stride, padding, dtype):
strides = (stride, stride)
shape_data = (N, CI, HI, WI)
shape_weight = (CO, CI, KSIZE, KSIZE)
shape_grad_output = (N, CO, HO, WO)
# given tensor
data = te.placeholder(shape_data, name="data", dtype=dtype)
weight = te.placeholder(shape_weight, name="weight", dtype=dtype)
grad_output = te.placeholder(shape_grad_output,
name="grad_output",
dtype=dtype)
# grad_data
out_h = (HO - 1) * strides[0] - 2 * padding + KSIZE
out_w = (WO - 1) * strides[1] - 2 * padding + KSIZE
output_padding = (HI - out_h, WI - out_w)
grad_data = topi.nn.conv2d_transpose_nchw(grad_output, weight, strides,
padding, dtype, output_padding)
# grad_weight
dilation_h, dilation_w = (1, 1)
batch, in_channel, in_h, in_w = shape_data
out_channel, _, filter_h, filter_w = shape_weight
grad_output_tmp = topi.tile(grad_output, [1, in_channel, 1, 1])
grad_output_tmp = topi.reshape(
grad_output_tmp, [batch * in_channel * out_channel, 1, HO, WO])
data_tmp = topi.reshape(data, [1, in_channel * batch, HI, WI])
grad_weight = topi.nn.group_conv2d_nchw(data_tmp,
grad_output_tmp,
stride=(dilation_h, dilation_w),
padding=padding,
dilation=strides,
groups=in_channel * batch,
out_dtype=dtype)
# infer shape of grad_weight
_, _, grad_h, grad_w = shape_grad_output
fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(
padding, (filter_h, filter_w))
padded_weight_grad_h = (in_h - (grad_h - 1) * strides[0] - 1 + fpad_top +
fpad_bottom) // dilation_h + 1
padded_weight_grad_w = (in_w - (grad_w - 1) * strides[1] - 1 + fpad_left +
fpad_right) // dilation_w + 1
grad_weight = topi.reshape(grad_weight, [
batch, in_channel, out_channel, padded_weight_grad_h,
padded_weight_grad_w
])
grad_weight = topi.sum(grad_weight, axis=0)
grad_weight = topi.transpose(grad_weight, [1, 0, 2, 3])
if padded_weight_grad_h > filter_h or padded_weight_grad_w > filter_w:
grad_weight = topi.strided_slice(
grad_weight,
begin=[0, 0, 0, 0],
end=[out_channel, in_channel, filter_h, filter_w])
return [data, weight, grad_output, grad_data, grad_weight]
return [data, weight, grad_output, grad_data, grad_weight]