setup.py 1.75 KB
Newer Older
anton's avatar
anton committed
1
2
3
#!/usr/bin/env python
import os
from setuptools import setup, find_packages
anton's avatar
anton committed
4
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
5

anton's avatar
anton committed
6
7
8
9
10

with open('requirements.txt') as f:
    requirements = f.read().splitlines()


anton's avatar
anton committed
11
long_description = """
Daniel Povey's avatar
Daniel Povey committed
12
13
This package implements an efficient parallel algorithm for the computation of
mutual information between sequences with differentiable bindings to PyTorch.
anton's avatar
anton committed
14
15
16


Find more details and the most up-to-date information on the project webpage:
Daniel Povey's avatar
Daniel Povey committed
17
[TODO]
anton's avatar
anton committed
18
19
20
"""


21
22
23
def configure_extensions():
    out = [
        CppExtension(
Daniel Povey's avatar
Daniel Povey committed
24
            'torch_mutual_information_cpu',
25
            [
Daniel Povey's avatar
Daniel Povey committed
26
                os.path.join('torch_mutual_information', 'mutual_information_cpu.cpp'),
27
28
29
30
31
32
            ],
        )
    ]
    try:
        out.append(
            CUDAExtension(
Daniel Povey's avatar
Daniel Povey committed
33
                'torch_mutual_information_cuda',
34
                [
Daniel Povey's avatar
Daniel Povey committed
35
36
                    os.path.join('torch_mutual_information', 'mutual_information_cuda.cpp'),
                    os.path.join('torch_mutual_information', 'mutual_information_cuda_kernel.cu'),
37
38
39
40
41
42
43
44
                ],
            )
        )
    except Exception as e:
        print(f'Failed to build CUDA extension, this part of the package will not work. Reason: {str(e)}')
    return out


45
setup(
Daniel Povey's avatar
Daniel Povey committed
46
    name='torch_mutual_information',
anton's avatar
anton committed
47
    version='1.0.2',
Daniel Povey's avatar
Daniel Povey committed
48
    description='Mutual information between sequences of vectors',
anton's avatar
anton committed
49
50
    long_description=long_description,
    long_description_content_type='text/markdown',
anton's avatar
anton committed
51
52
53
    install_requires=requirements,
    python_requires='>=3.6',
    packages=find_packages(),
54
    author='Dan Povey',
anton's avatar
anton committed
55
    license='BSD',
56
    ext_modules=configure_extensions(),
57
58
    cmdclass={
        'build_ext': BuildExtension
anton's avatar
anton committed
59
60
    },
    keywords=[
Daniel Povey's avatar
Daniel Povey committed
61
        'pytorch', 'sequence', 'mutual', 'information'
anton's avatar
anton committed
62
63
    ],
)