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