| import contextlib
|
| import warnings
|
|
|
| import torch
|
| from torch import autograd
|
| from torch.nn import functional as F
|
|
|
| enabled = True
|
| weight_gradients_disabled = False
|
|
|
|
|
| @contextlib.contextmanager
|
| def no_weight_gradients():
|
| global weight_gradients_disabled
|
|
|
| old = weight_gradients_disabled
|
| weight_gradients_disabled = True
|
| yield
|
| weight_gradients_disabled = old
|
|
|
|
|
| def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
|
| if could_use_op(input):
|
| return conv2d_gradfix(
|
| transpose=False,
|
| weight_shape=weight.shape,
|
| stride=stride,
|
| padding=padding,
|
| output_padding=0,
|
| dilation=dilation,
|
| groups=groups,
|
| ).apply(input, weight, bias)
|
|
|
| return F.conv2d(
|
| input=input,
|
| weight=weight,
|
| bias=bias,
|
| stride=stride,
|
| padding=padding,
|
| dilation=dilation,
|
| groups=groups,
|
| )
|
|
|
|
|
| def conv_transpose2d(
|
| input,
|
| weight,
|
| bias=None,
|
| stride=1,
|
| padding=0,
|
| output_padding=0,
|
| groups=1,
|
| dilation=1,
|
| ):
|
| if could_use_op(input):
|
| return conv2d_gradfix(
|
| transpose=True,
|
| weight_shape=weight.shape,
|
| stride=stride,
|
| padding=padding,
|
| output_padding=output_padding,
|
| groups=groups,
|
| dilation=dilation,
|
| ).apply(input, weight, bias)
|
|
|
| return F.conv_transpose2d(
|
| input=input,
|
| weight=weight,
|
| bias=bias,
|
| stride=stride,
|
| padding=padding,
|
| output_padding=output_padding,
|
| dilation=dilation,
|
| groups=groups,
|
| )
|
|
|
|
|
| def could_use_op(input):
|
| if (not enabled) or (not torch.backends.cudnn.enabled):
|
| return False
|
|
|
| if input.device.type != "cuda":
|
| return False
|
|
|
| if any(torch.__version__.startswith(x) for x in ["1.7.", "1.8."]):
|
| return True
|
|
|
| warnings.warn(
|
| f"conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d()."
|
| )
|
|
|
| return False
|
|
|
|
|
| def ensure_tuple(xs, ndim):
|
| xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
|
|
|
| return xs
|
|
|
|
|
| conv2d_gradfix_cache = dict()
|
|
|
|
|
| def conv2d_gradfix(
|
| transpose, weight_shape, stride, padding, output_padding, dilation, groups
|
| ):
|
| ndim = 2
|
| weight_shape = tuple(weight_shape)
|
| stride = ensure_tuple(stride, ndim)
|
| padding = ensure_tuple(padding, ndim)
|
| output_padding = ensure_tuple(output_padding, ndim)
|
| dilation = ensure_tuple(dilation, ndim)
|
|
|
| key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
|
| if key in conv2d_gradfix_cache:
|
| return conv2d_gradfix_cache[key]
|
|
|
| common_kwargs = dict(
|
| stride=stride, padding=padding, dilation=dilation, groups=groups
|
| )
|
|
|
| def calc_output_padding(input_shape, output_shape):
|
| if transpose:
|
| return [0, 0]
|
|
|
| return [
|
| input_shape[i + 2]
|
| - (output_shape[i + 2] - 1) * stride[i]
|
| - (1 - 2 * padding[i])
|
| - dilation[i] * (weight_shape[i + 2] - 1)
|
| for i in range(ndim)
|
| ]
|
|
|
| class Conv2d(autograd.Function):
|
| @staticmethod
|
| def forward(ctx, input, weight, bias):
|
| if not transpose:
|
| out = F.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
|
|
|
| else:
|
| out = F.conv_transpose2d(
|
| input=input,
|
| weight=weight,
|
| bias=bias,
|
| output_padding=output_padding,
|
| **common_kwargs,
|
| )
|
|
|
| ctx.save_for_backward(input, weight)
|
|
|
| return out
|
|
|
| @staticmethod
|
| def backward(ctx, grad_output):
|
| input, weight = ctx.saved_tensors
|
| grad_input, grad_weight, grad_bias = None, None, None
|
|
|
| if ctx.needs_input_grad[0]:
|
| p = calc_output_padding(
|
| input_shape=input.shape, output_shape=grad_output.shape
|
| )
|
| grad_input = conv2d_gradfix(
|
| transpose=(not transpose),
|
| weight_shape=weight_shape,
|
| output_padding=p,
|
| **common_kwargs,
|
| ).apply(grad_output, weight, None)
|
|
|
| if ctx.needs_input_grad[1] and not weight_gradients_disabled:
|
| grad_weight = Conv2dGradWeight.apply(grad_output, input)
|
|
|
| if ctx.needs_input_grad[2]:
|
| grad_bias = grad_output.sum((0, 2, 3))
|
|
|
| return grad_input, grad_weight, grad_bias
|
|
|
| class Conv2dGradWeight(autograd.Function):
|
| @staticmethod
|
| def forward(ctx, grad_output, input):
|
| op = torch._C._jit_get_operation(
|
| "aten::cudnn_convolution_backward_weight"
|
| if not transpose
|
| else "aten::cudnn_convolution_transpose_backward_weight"
|
| )
|
| flags = [
|
| torch.backends.cudnn.benchmark,
|
| torch.backends.cudnn.deterministic,
|
| torch.backends.cudnn.allow_tf32,
|
| ]
|
| grad_weight = op(
|
| weight_shape,
|
| grad_output,
|
| input,
|
| padding,
|
| stride,
|
| dilation,
|
| groups,
|
| *flags,
|
| )
|
| ctx.save_for_backward(grad_output, input)
|
|
|
| return grad_weight
|
|
|
| @staticmethod
|
| def backward(ctx, grad_grad_weight):
|
| grad_output, input = ctx.saved_tensors
|
| grad_grad_output, grad_grad_input = None, None
|
|
|
| if ctx.needs_input_grad[0]:
|
| grad_grad_output = Conv2d.apply(input, grad_grad_weight, None)
|
|
|
| if ctx.needs_input_grad[1]:
|
| p = calc_output_padding(
|
| input_shape=input.shape, output_shape=grad_output.shape
|
| )
|
| grad_grad_input = conv2d_gradfix(
|
| transpose=(not transpose),
|
| weight_shape=weight_shape,
|
| output_padding=p,
|
| **common_kwargs,
|
| ).apply(grad_output, grad_grad_weight, None)
|
|
|
| return grad_grad_output, grad_grad_input
|
|
|
| conv2d_gradfix_cache[key] = Conv2d
|
|
|
| return Conv2d
|
|
|
