DL-Art-School/codes/models/archs/dcn/setup.py
XintaoWang 037933ba66 mmsr
2019-08-23 21:42:47 +08:00

23 lines
711 B
Python

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
def make_cuda_ext(name, sources):
return CUDAExtension(
name='{}'.format(name), sources=[p for p in sources], extra_compile_args={
'cxx': [],
'nvcc': [
'-D__CUDA_NO_HALF_OPERATORS__',
'-D__CUDA_NO_HALF_CONVERSIONS__',
'-D__CUDA_NO_HALF2_OPERATORS__',
]
})
setup(
name='deform_conv', ext_modules=[
make_cuda_ext(name='deform_conv_cuda',
sources=['src/deform_conv_cuda.cpp', 'src/deform_conv_cuda_kernel.cu'])
], cmdclass={'build_ext': BuildExtension}, zip_safe=False)