cupy_convolution_forward_example.py

dtype = numpy.float32
batch_size = 2
in_channels = 3
out_channels =2
kernel_size = (3, 3, 1)
stride = (1, 1, 1)
pad = (0, 0, 0)
in_dims = (1000, 768, 3)
out_dims = (998, 766, 3)
x_shape = (batch_size, in_channels) + in_dims
w_shape = (out_channels, in_channels) + kernel_size
b_shape = (out_channels,)
y_shape = (batch_size, out_channels) + out_dims

def total_size(shape):
    return functools.reduce(operator.mul, shape, 1)

x = cupy.arange(start=-total_size(x_shape)//2, stop=total_size(x_shape)//2, dtype=dtype).reshape(x_shape)
W = cupy.arange(start=-total_size(w_shape)//2, stop=total_size(w_shape)//2, dtype=dtype).reshape(w_shape)
b = cupy.array([-0.2, 1.3], dtype=dtype)
y = cupy.empty(y_shape, dtype=dtype)

cudnn.convolution_forward(
    x, W, b, y,
    pad, stride, None, 1,
    auto_tune=True, tensor_core='never')