forked from NateMeyer/tensorrt_demos
-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup.py
47 lines (42 loc) · 1.15 KB
/
setup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
from Cython.Build import cythonize
import numpy
library_dirs = [
'/usr/local/cuda/lib64',
'/usr/local/TensorRT-7.1.3.4/lib', # for my x86_64 PC
'/usr/local/lib',
]
libraries = [
'nvinfer',
'cudnn',
'cublas',
'cudart_static',
'nvToolsExt',
'cudart',
'rt',
]
include_dirs = [
# in case the following numpy include path does not work, you
# could replace it manually with, say,
# '-I/usr/local/lib/python3.6/dist-packages/numpy/core/include',
'-I' + numpy.__path__[0] + '/core/include',
'-I/usr/local/cuda/include',
'-I/usr/local/TensorRT-7.1.3.4/include', # for my x86_64 PC
'-I/usr/local/include',
]
setup(
cmdclass={'build_ext': build_ext},
ext_modules=cythonize(
Extension(
'pytrt',
sources=['pytrt.pyx'],
language='c++',
library_dirs=library_dirs,
libraries=libraries,
extra_compile_args=['-O3', '-std=c++11'] + include_dirs
),
compiler_directives={'language_level': '3'}
)
)