Update vtoonify/model/stylegan/op/conv2d_gradfix.py

vtoonify/model/stylegan/op/conv2d_gradfix.py

@@ -1,227 +1,227 @@
-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
+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