Skip to content

Commit 104e21e

Browse files
zmiczmic
committed
xxx
1 parent 69934d5 commit 104e21e

File tree

1 file changed

+168
-0
lines changed

1 file changed

+168
-0
lines changed

‎lib/nms/setup_windows.py

+168
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
# --------------------------------------------------------
2+
# Pose.gluon
3+
# Copyright (c) 2018-present Microsoft
4+
# Licensed under The MIT License [see LICENSE for details]
5+
# Modified from py-faster-rcnn (https://github.com/rbgirshick/py-faster-rcnn)
6+
# --------------------------------------------------------
7+
8+
import os
9+
from os.path import join as pjoin
10+
from setuptools import setup
11+
from distutils.extension import Extension
12+
from Cython.Distutils import build_ext
13+
import numpy as np
14+
from shutil import which
15+
16+
17+
def find_in_path(name, path):
18+
"Find a file in a search path"
19+
# Adapted fom
20+
# http://code.activestate.com/recipes/52224-find-a-file-given-a-search-path/
21+
for dir in path.split(os.pathsep):
22+
binpath = pjoin(dir, name)
23+
if os.path.exists(binpath):
24+
return os.path.abspath(binpath)
25+
return None
26+
27+
28+
def locate_cuda():
29+
"""Locate the CUDA environment on the system
30+
Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64'
31+
and values giving the absolute path to each directory.
32+
Starts by looking for the CUDA_PATH env variable. If not found, everything
33+
is based on finding 'nvcc.exe' in the PATH.
34+
"""
35+
36+
# CUDA_PATH C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2
37+
if 'CUDA_PATH1' in os.environ:
38+
home = os.environ['CUDA_PATH']
39+
nvcc = pjoin(home, 'bin', 'nvcc.exe')
40+
else:
41+
nvcc = which('nvcc.exe')
42+
if nvcc is None:
43+
raise EnvironmentError('The nvcc binary could not be '
44+
'located in your $PATH. Either add it to your path, or set $CUDA_PATH')
45+
home = os.path.dirname(os.path.dirname(nvcc))
46+
47+
cudaconfig = {'home':home, 'nvcc':nvcc,
48+
'include': pjoin(home, 'include'),
49+
'lib64': pjoin(home, 'lib', 'x64')}
50+
for k, v in cudaconfig.items():
51+
if not os.path.exists(v):
52+
raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v))
53+
54+
return cudaconfig
55+
CUDA = locate_cuda()
56+
57+
58+
# Obtain the numpy include directory. This logic works across numpy versions.
59+
try:
60+
numpy_include = np.get_include()
61+
except AttributeError:
62+
numpy_include = np.get_numpy_include()
63+
64+
65+
def customize_compiler_for_nvcc(self):
66+
"""inject deep into distutils to customize how the dispatch
67+
to gcc/nvcc works.
68+
If you subclass UnixCCompiler, it's not trivial to get your subclass
69+
injected in, and still have the right customizations (i.e.
70+
distutils.sysconfig.customize_compiler) run on it. So instead of going
71+
the OO route, I have this. Note, it's kindof like a wierd functional
72+
subclassing going on."""
73+
74+
#print(self.__class__.__dict__)
75+
#print(self.src_extensions)
76+
77+
# tell the compiler it can processes .cu
78+
self.src_extensions.append('.cu')
79+
#self.set_executable('compiler', CUDA['nvcc'])
80+
81+
# save references to the default compiler_so and _comple methods
82+
#default_compiler_so = self.compiler_so
83+
super = self.compile
84+
85+
def compile(sources,
86+
output_dir=None, macros=None, include_dirs=None, debug=0,
87+
extra_preargs=None, extra_postargs=None, depends=None):
88+
sources_cpp = []
89+
for src in sources:
90+
if os.path.splitext(src)[1] == '.cu':
91+
# use the cuda for .cu files
92+
args = [CUDA['nvcc']] + extra_postargs['nvcc'] + [src]
93+
print(args)
94+
if not self.initialized:
95+
self.initialize()
96+
compile_info = self._setup_compile(output_dir, macros, include_dirs,
97+
sources, depends, extra_postargs)
98+
macros, objects, extra_postargs, pp_opts, build = compile_info
99+
self.spawn(args)
100+
else:
101+
sources_cpp.append(src)
102+
103+
super(sources_cpp,
104+
output_dir, macros, include_dirs, debug,
105+
extra_preargs, extra_postargs['cl'], depends)
106+
107+
108+
# now redefine the _compile method. This gets executed for each
109+
# object but distutils doesn't have the ability to change compilers
110+
# based on source extension: we add it.
111+
def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
112+
if os.path.splitext(src)[1] == '.cu':
113+
# use the cuda for .cu files
114+
print("!", src)
115+
self.set_executable('compiler', CUDA['nvcc'])
116+
# use only a subset of the extra_postargs, which are 1-1 translated
117+
# from the extra_compile_args in the Extension class
118+
postargs = extra_postargs['nvcc']
119+
else:
120+
postargs = extra_postargs['cl']
121+
122+
super(obj, src, ext, cc_args, postargs, pp_opts)
123+
# reset the default compiler_so, which we might have changed for cuda
124+
#self.compiler_so = default_compiler_so
125+
126+
# inject our redefined _compile method into the class
127+
self.compile = compile
128+
129+
130+
# run the customize_compiler
131+
class custom_build_ext(build_ext):
132+
def build_extensions(self):
133+
customize_compiler_for_nvcc(self.compiler)
134+
build_ext.build_extensions(self)
135+
136+
137+
ext_modules = [
138+
Extension(
139+
"cpu_nms",
140+
["cpu_nms.pyx"],
141+
extra_compile_args={'cl': []},
142+
include_dirs = [numpy_include]
143+
),
144+
Extension('gpu_nms',
145+
['nms_kernel.cu', 'gpu_nms.pyx'],
146+
library_dirs=[CUDA['lib64']],
147+
libraries=['cudart'],
148+
language='c++',
149+
runtime_library_dirs=[CUDA['lib64']],
150+
# this syntax is specific to this build system
151+
# we're only going to use certain compiler args with nvcc and not with
152+
# gcc the implementation of this trick is in customize_compiler() below
153+
extra_compile_args={'cl': [],
154+
'nvcc': ['-arch=sm_35',
155+
'--ptxas-options=-v',
156+
'-c',
157+
'--compiler-options',
158+
"'-fPIC'"]},
159+
include_dirs = [numpy_include, CUDA['include']]
160+
),
161+
]
162+
163+
setup(
164+
name='nms',
165+
ext_modules=ext_modules,
166+
# inject our custom trigger
167+
cmdclass={'build_ext': custom_build_ext},
168+
)

0 commit comments

Comments
 (0)