def forward(ctx,
x,
w,
bias=None,
stride=(1, 1),
padding=(0, 0),
output_padding=(0, 0),
dilation=(1, 1),
groups=1):
# Save constants for bprop
ctx.stride = stride
ctx.padding = padding
ctx.output_padding = output_padding
ctx.dilation = dilation
ctx.groups = groups
ctx.need_bias_grad = bias is not None
ctx.save_for_backward(x, w)
if bias is None:
return NHWC.cudnn_convolution_transpose_nhwc(
x, w, padding, output_padding, stride, dilation, groups,
torch.backends.cudnn.benchmark, False)
else:
return NHWC.cudnn_convolution_transpose_with_bias_nhwc(
x, w, bias, padding, output_padding, stride, dilation, groups,
torch.backends.cudnn.benchmark, False)
def backward(ctx, grad_y):
x, w = ctx.saved_variables
## if padding is used in fprop, we should pad grad_y here also for better perf
K = grad_y.shape[3]
is_padded = (K % 8) != 0
if is_padded:
K_padded = 8 * ((K + 7) // 8)
padded_grad_shape = [grad_y.shape[0], grad_y.shape[1], grad_y.shape[2], K_padded]
padded_grad = torch.zeros(padded_grad_shape, dtype = grad_y.dtype, device = grad_y.device)
padded_grad[:,:,:,:K] = grad_y
# print("padded grad shape", padded_grad.shape)
#grad_y = padded_grad
if ctx.need_bias_grad:
if not is_padded:
dx, dw, db = NHWC.cudnn_convolution_backward_with_bias_nhwc(x, grad_y, w,
ctx.padding, ctx.stride, ctx.dilation, ctx.groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic,
list(ctx.needs_input_grad[0:3]))
else:
dx, dw, db = NHWC.cudnn_convolution_backward_with_bias_nhwc(x, padded_grad, w,
ctx.padding, ctx.stride, ctx.dilation, ctx.groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic,
list(ctx.needs_input_grad[0:3]))
if ctx.needs_input_grad[0]:
if not is_padded:
return dx, dw, db, None, None, None, None
else:
return dx, dw[:K,:,:,:].contiguous(), db[:K].contiguous(), None, None, None, None
else:
if not is_padded:
return None, dw, db, None, None, None, None
else:
return None, dw[:K,:,:,:].contiguous(), db[:K].contiguous(), None, None, None, None
else:
if (not ctx.needs_input_grad[1] ):
return None, None, None, None, None, None, None
if not is_padded:
dx, dw = NHWC.cudnn_convolution_backward_nhwc(x, grad_y, w,
ctx.padding, ctx.stride, ctx.dilation, ctx.groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic,
list(ctx.needs_input_grad[0:2]))
else:
dx, dw = NHWC.cudnn_convolution_backward_nhwc(x, padded_grad, w,
ctx.padding, ctx.stride, ctx.dilation, ctx.groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic,
list(ctx.needs_input_grad[0:2]))
if (not ctx.needs_input_grad[1]):
return None, None, None, None, None, None, None
elif ctx.needs_input_grad[0]:
if not is_padded:
return dx, dw, None, None, None, None, None
else:
return dx, dw[:K,:,:,:].contiguous(), None, None, None, None, None
else:
if not is_padded:
return None, dw, None, None, None, None, None
else:
return None, dw[:K,:,:,:].contiguous(), None, None, None, None, None
def forward(ctx, x, w, bias=None, stride=(1,1), padding=(0,0), dilation=(1,1), groups=1):
# Save constants for bprop
ctx.stride = stride
ctx.padding = padding
ctx.dilation = dilation
ctx.groups = groups
ctx.need_bias_grad = bias is not None
## pad K-dimension to multiple of 8 for better perf
K = w.shape[0]
is_padded = (K % 8) != 0
if is_padded:
K_padded = 8 * ((K + 7) // 8)
padded_filter_shape = [K_padded, w.shape[1], w.shape[2], w.shape[3]]
padded_w = torch.zeros(padded_filter_shape, dtype = w.dtype, device = w.device)
padded_w[:K,:,:,:] = w
ctx.save_for_backward(x, padded_w)##debug##, only use padding in fprop
# ctx.save_for_backward(x, w)
if bias is not None:
padded_bias = torch.zeros([K_padded], dtype = bias.dtype, device = bias.device)
padded_bias[:K] = bias
else:
ctx.save_for_backward(x, w)
# try padding only in grad
if bias is None:
if is_padded:
output = NHWC.cudnn_convolution_nhwc(x, padded_w,
padding, stride, dilation,
groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic)
return output[:,:,:,:K].contiguous()
else:
return NHWC.cudnn_convolution_nhwc(x, w,
padding, stride, dilation,
groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic)
else:
if is_padded:
output = NHWC.cudnn_convolution_with_bias_nhwc(x, padded_w, padded_bias,
padding, stride, dilation,
groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic)
return output[:,:,:,:K].contiguous()
else:
return NHWC.cudnn_convolution_with_bias_nhwc(x, w, bias,
padding, stride, dilation,
groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic)
def backward(ctx, y_grad):
x, y = ctx.saved_variables
kernel = ctx.kernel_size
stride = ctx.stride
padding = ctx.padding
dilation = ctx.dilation
return NHWC.max_pool_bwd_nhwc(x, y, y_grad, kernel, stride, padding,
dilation), None, None, None, None
def forward(ctx, x, kernel_size, stride, padding, dilation):
ctx.kernel_size = kernel_size
ctx.stride = stride if stride is not None else 0
ctx.padding = padding
ctx.dilation = dilation
y = NHWC.max_pool_fwd_nhwc(x, kernel_size, stride, padding, dilation)
# Need to save y as well
ctx.save_for_backward(x, y)
return y
def backward(ctx, grad_y):
x, w = ctx.saved_variables
if ctx.need_bias_grad:
dx, dw, db = NHWC.cudnn_convolution_transpose_backward_with_bias_nhwc(x, grad_y, w,
ctx.padding, ctx.output_padding, ctx.stride, ctx.dilation, ctx.groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic,
list(ctx.needs_input_grad[0:3]))
if ctx.needs_input_grad[0]:
return dx, dw, db, None, None, None, None, None
else:
return None, dw, db, None, None, None, None, None
else:
if (not ctx.needs_input_grad[1] and not ctx.needs_input_grad[0]):
return None, None, None, None, None, None, None
dx, dw = NHWC.cudnn_convolution_transpose_backward_nhwc(x, grad_y, w,
ctx.padding, ctx.output_padding, ctx.stride, ctx.dilation, ctx.groups,
torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic,
list(ctx.needs_input_grad[0:2]))
if (not ctx.needs_input_grad[1]):
return None, None, None, None, None, None, None, None
elif ctx.needs_input_grad[0]:
return dx, dw, None, None, None, None, None, None
else:
return None, dw, None, None, None, None, None, None
def backward(ctx, y_grad):
input_size = ctx.input_size
output_size = ctx.output_size
return NHWC.upsample_nearest2d_backward_cuda(y_grad, output_size,
input_size), None
def forward(ctx, x, output_size):
ctx.output_size = output_size
ctx.input_size = x.shape
y = NHWC.upsample_nearest2d_cuda(x, output_size)
# Need to save y as well
return y
def backward(ctx, y_grad):
x_grad = NHWC.cudnnNhwcToNchw(y_grad)
return x_grad
def forward(ctx, x):
y = NHWC.cudnnNchwToNhwc(x)
return y
def forward(ctx, x):
y = NHWC.cudnnNhwcToNchw(x)
return y