Print better error message for the dreaded upfirdn2d_plugin problem

Print full traceback when custom extension build fails.

Also allow pytorch 1.9 so that this runs against pytorch upstream
devel builds.

issues #2, #28, #35, #37, #39

Author: jannehellsten

torch_utils/ops/bias_act.py

@@ -9,11 +9,11 @@
 """Custom PyTorch ops for efficient bias and activation."""
 
 import os
-import sys
 import warnings
 import numpy as np
 import torch
 import dnnlib
+import traceback
 
 from .. import custom_ops
 from .. import misc
@@ -47,7 +47,7 @@ def _init():
         try:
             _plugin = custom_ops.get_plugin('bias_act_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
         except:
-            warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + str(sys.exc_info()[1]))
+            warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
     return _plugin is not None
 
 #----------------------------------------------------------------------------

torch_utils/ops/conv2d_gradfix.py

@@ -50,7 +50,7 @@ def _should_use_custom_op(input):
         return False
     if input.device.type != 'cuda':
         return False
-    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.']):
+    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
         return True
     warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().')
     return False

torch_utils/ops/grid_sample_gradfix.py

@@ -34,7 +34,7 @@ def grid_sample(input, grid):
 def _should_use_custom_op():
     if not enabled:
         return False
-    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.']):
+    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
         return True
     warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')
     return False

torch_utils/ops/upfirdn2d.py

@@ -9,10 +9,10 @@
 """Custom PyTorch ops for efficient resampling of 2D images."""
 
 import os
-import sys
 import warnings
 import numpy as np
 import torch
+import traceback
 
 from .. import custom_ops
 from .. import misc
@@ -31,7 +31,7 @@ def _init():
         try:
             _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
         except:
-            warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + str(sys.exc_info()[1]))
+            warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
     return _plugin is not None
 
 def _parse_scaling(scaling):