Janne Hellsten commited on
Commit
2506395
1 Parent(s): 386669a

Print better error message for the dreaded upfirdn2d_plugin problem

Browse files

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

torch_utils/ops/bias_act.py CHANGED
@@ -9,11 +9,11 @@
9
  """Custom PyTorch ops for efficient bias and activation."""
10
 
11
  import os
12
- import sys
13
  import warnings
14
  import numpy as np
15
  import torch
16
  import dnnlib
 
17
 
18
  from .. import custom_ops
19
  from .. import misc
@@ -47,7 +47,7 @@ def _init():
47
  try:
48
  _plugin = custom_ops.get_plugin('bias_act_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
49
  except:
50
- warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + str(sys.exc_info()[1]))
51
  return _plugin is not None
52
 
53
  #----------------------------------------------------------------------------
 
9
  """Custom PyTorch ops for efficient bias and activation."""
10
 
11
  import os
 
12
  import warnings
13
  import numpy as np
14
  import torch
15
  import dnnlib
16
+ import traceback
17
 
18
  from .. import custom_ops
19
  from .. import misc
 
47
  try:
48
  _plugin = custom_ops.get_plugin('bias_act_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
49
  except:
50
+ warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
51
  return _plugin is not None
52
 
53
  #----------------------------------------------------------------------------
torch_utils/ops/conv2d_gradfix.py CHANGED
@@ -50,7 +50,7 @@ def _should_use_custom_op(input):
50
  return False
51
  if input.device.type != 'cuda':
52
  return False
53
- if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.']):
54
  return True
55
  warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().')
56
  return False
 
50
  return False
51
  if input.device.type != 'cuda':
52
  return False
53
+ if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
54
  return True
55
  warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().')
56
  return False
torch_utils/ops/grid_sample_gradfix.py CHANGED
@@ -34,7 +34,7 @@ def grid_sample(input, grid):
34
  def _should_use_custom_op():
35
  if not enabled:
36
  return False
37
- if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.']):
38
  return True
39
  warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')
40
  return False
 
34
  def _should_use_custom_op():
35
  if not enabled:
36
  return False
37
+ if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
38
  return True
39
  warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')
40
  return False
torch_utils/ops/upfirdn2d.py CHANGED
@@ -9,10 +9,10 @@
9
  """Custom PyTorch ops for efficient resampling of 2D images."""
10
 
11
  import os
12
- import sys
13
  import warnings
14
  import numpy as np
15
  import torch
 
16
 
17
  from .. import custom_ops
18
  from .. import misc
@@ -31,7 +31,7 @@ def _init():
31
  try:
32
  _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
33
  except:
34
- warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + str(sys.exc_info()[1]))
35
  return _plugin is not None
36
 
37
  def _parse_scaling(scaling):
 
9
  """Custom PyTorch ops for efficient resampling of 2D images."""
10
 
11
  import os
 
12
  import warnings
13
  import numpy as np
14
  import torch
15
+ import traceback
16
 
17
  from .. import custom_ops
18
  from .. import misc
 
31
  try:
32
  _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
33
  except:
34
+ warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
35
  return _plugin is not None
36
 
37
  def _parse_scaling(scaling):