[Fix] Enhance the compatibility of training stylegan 2 (#2694)

* [Fix] Enhance the compatibility of training stylegan 2

* Fix unit test in PyTorch 2.0

* Apply suggestions from code review

---------
Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

mmcv/ops/conv2d_gradfix.py

@@ -15,6 +15,7 @@ import warnings
 from typing import Dict, Optional, Tuple, Union
 
 import torch
+from mmengine.utils import digit_version
 
 enabled = True
 weight_gradients_disabled = False
@@ -259,17 +260,43 @@ def _conv2d_gradfix(
                     memory_format=(torch.channels_last if input.stride(1) ==
                                    1 else torch.contiguous_format))
 
-            # General case => cuDNN.
-            name = ('aten::cudnn_convolution_transpose_backward_weight' if
-                    transpose else 'aten::cudnn_convolution_backward_weight')
-            flags = [
-                torch.backends.cudnn.benchmark,
-                torch.backends.cudnn.deterministic,
-                torch.backends.cudnn.allow_tf32
-            ]
-            return torch._C._jit_get_operation(name)(weight_shape, grad_output,
-                                                     input, padding, stride,
-                                                     dilation, groups, *flags)
+            # PyTorch consolidated convolution backward API in PR:
+            # https://github.com/pytorch/pytorch/commit/3dc3651e0ee3623f669c3a2c096408dbc476d122  # noqa: E501
+            # Enhance the code referring to the discussion:
+            # https://github.com/pytorch/pytorch/issues/74437
+            if digit_version(torch.__version__) >= digit_version('1.11.0'):
+                empty_weight = torch.tensor(
+                    0.0, dtype=input.dtype,
+                    device=input.device).expand(weight_shape)
+                output_padding = calc_output_padding(input.shape,
+                                                     grad_output.shape)
+                return torch.ops.aten.convolution_backward(
+                    grad_output,
+                    input,
+                    empty_weight,
+                    None,
+                    stride=stride,
+                    dilation=dilation,
+                    transposed=transpose,
+                    padding=padding,
+                    groups=groups,
+                    output_padding=output_padding,
+                    output_mask=[0, 1, 0])[1]
+            else:
+                # General case => cuDNN.
+                name = ('aten::cudnn_convolution_transpose_backward_weight'
+                        if transpose else
+                        'aten::cudnn_convolution_backward_weight')
+                flags = [
+                    torch.backends.cudnn.benchmark,
+                    torch.backends.cudnn.deterministic,
+                    torch.backends.cudnn.allow_tf32
+                ]
+                return torch._C._jit_get_operation(name)(weight_shape,
+                                                         grad_output, input,
+                                                         padding, stride,
+                                                         dilation, groups,
+                                                         *flags)
 
         @staticmethod
         def backward(ctx, grad2_grad_weight):

tests/test_ops/test_conv_gradfix.py

@@ -20,8 +20,8 @@ class TestCond2d:
         weight = self.weight.cuda()
         res = conv2d(x, weight, None, 1, 1)
         assert res.shape == (1, 1, 32, 32)
-        gradcheck(conv2d, (x, weight, None, 1, 1), eps=1e-2, atol=1e-2)
-        gradgradcheck(conv2d, (x, weight, None, 1, 1), eps=1e-2, atol=1e-2)
+        gradcheck(conv2d, (x, weight, None, 1, 1), eps=1e-2, atol=0.1)
+        gradgradcheck(conv2d, (x, weight, None, 1, 1), eps=1e-2, atol=0.1)
 
 
 class TestCond2dTansposed:
@@ -37,7 +37,6 @@ class TestCond2dTansposed:
         weight = self.weight.cuda()
         res = conv_transpose2d(x, weight, None, 1, 1)
         assert res.shape == (1, 1, 32, 32)
-        gradcheck(
-            conv_transpose2d, (x, weight, None, 1, 1), eps=1e-2, atol=1e-2)
+        gradcheck(conv_transpose2d, (x, weight, None, 1, 1), eps=1e-2, atol=1e-2)
         gradgradcheck(
             conv_transpose2d, (x, weight, None, 1, 1), eps=1e-2, atol=1e-2)