Commit 25063950 authored by Janne Hellsten's avatar Janne Hellsten
Browse files

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
parent 386669a8

torch_utils/ops/bias_act.py

+2 −2
Original line number Diff line number Diff line
@@ -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

+1 −1
Original line number Diff line number Diff line
@@ -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

+1 −1
Original line number Diff line number Diff line
@@ -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

+2 −2
Original line number Diff line number Diff line
@@ -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):