def forward(ctx, input, kernel=2, stride=None):
no_batch = False
if len(input.size()) == 3:
no_batch = True
input.unsqueeze_(0)
B, C, H, W = input.size()
kernel = _pair(kernel)
if stride is None:
stride = kernel
else:
stride = _pair(stride)
oH = H // stride[0]
oW = W // stride[1]
output = input.new_zeros((B, C, oH, oW))
softpool_cuda.forward_2d(input, kernel, stride, output)
ctx.save_for_backward(input)
ctx.kernel = kernel
ctx.stride = stride
if no_batch:
return output.squeeze_(0)
return output
def forward(ctx, input, kernel=2, stride=None):
# Create contiguous tensor (if tensor is not contiguous)
no_batch = False
if len(input.size()) == 3:
no_batch = True
input.unsqueeze_(0)
B, C, H, W = input.size()
kernel = _pair(kernel)
if stride is None:
stride = kernel
else:
stride = _pair(stride)
oH = (H - kernel[0]) // stride[0] + 1
oW = (W - kernel[1]) // stride[1] + 1
output = input.new_zeros((B, C, oH, oW))
softpool_cuda.forward_2d(input.contiguous(), kernel, stride, output)
ctx.save_for_backward(input)
ctx.kernel = kernel
ctx.stride = stride
if no_batch:
return output.squeeze_(0)
return output