setup.py 2.06 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
from setuptools import setup, find_packages
Dmitry Leonov's avatar
Dmitry Leonov committed
2
from sys import argv
rusty1s's avatar
rusty1s committed
3
import torch
4
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
rusty1s's avatar
rusty1s committed
5

rusty1s's avatar
rusty1s committed
6
7
8
9
10
11
12
TORCH_MAJOR = int(torch.__version__.split('.')[0])
TORCH_MINOR = int(torch.__version__.split('.')[1])

extra_compile_args = []
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
    extra_compile_args += ['-DVERSION_GE_1_3']

rusty1s's avatar
rusty1s committed
13
ext_modules = [
rusty1s's avatar
rusty1s committed
14
15
16
17
    CppExtension('torch_spline_conv.basis_cpu', ['cpu/basis.cpp'],
                 extra_compile_args=extra_compile_args),
    CppExtension('torch_spline_conv.weighting_cpu', ['cpu/weighting.cpp'],
                 extra_compile_args=extra_compile_args),
rusty1s's avatar
rusty1s committed
18
19
]
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
rusty1s's avatar
rusty1s committed
20

Dmitry Leonov's avatar
Dmitry Leonov committed
21
22
23
24
25
26
27
GPU = True
for arg in argv:
    if arg == '--cpu':
        GPU = False
        argv.remove(arg)

if CUDA_HOME is not None and GPU:
rusty1s's avatar
rusty1s committed
28
    ext_modules += [
29
        CUDAExtension('torch_spline_conv.basis_cuda',
rusty1s's avatar
rusty1s committed
30
31
                      ['cuda/basis.cpp', 'cuda/basis_kernel.cu'],
                      extra_compile_args=extra_compile_args),
32
        CUDAExtension('torch_spline_conv.weighting_cuda',
rusty1s's avatar
rusty1s committed
33
34
                      ['cuda/weighting.cpp', 'cuda/weighting_kernel.cu'],
                      extra_compile_args=extra_compile_args),
rusty1s's avatar
rusty1s committed
35
36
    ]

rusty1s's avatar
rusty1s committed
37
__version__ = '1.1.1'
rusty1s's avatar
rusty1s committed
38
url = 'https://github.com/rusty1s/pytorch_spline_conv'
rusty1s's avatar
rusty1s committed
39

rusty1s's avatar
rusty1s committed
40
41
install_requires = []
setup_requires = ['pytest-runner']
rusty1s's avatar
rusty1s committed
42
43
44
45
tests_require = ['pytest', 'pytest-cov']

setup(
    name='torch_spline_conv',
rusty1s's avatar
rusty1s committed
46
    version=__version__,
rusty1s's avatar
typo  
rusty1s committed
47
    description=('Implementation of the Spline-Based Convolution Operator of '
48
                 'SplineCNN in PyTorch'),
rusty1s's avatar
rusty1s committed
49
50
    author='Matthias Fey',
    author_email='matthias.fey@tu-dortmund.de',
rusty1s's avatar
rusty1s committed
51
52
53
    url=url,
    download_url='{}/archive/{}.tar.gz'.format(url, __version__),
    keywords=[
rusty1s's avatar
rusty1s committed
54
55
56
57
        'pytorch',
        'geometric-deep-learning',
        'graph-neural-networks',
        'spline-cnn',
rusty1s's avatar
rusty1s committed
58
    ],
rusty1s's avatar
rusty1s committed
59
60
61
    install_requires=install_requires,
    setup_requires=setup_requires,
    tests_require=tests_require,
rusty1s's avatar
rusty1s committed
62
63
64
    ext_modules=ext_modules,
    cmdclass=cmdclass,
    packages=find_packages(),
rusty1s's avatar
rusty1s committed
65
)