From f3768bcdcf70477bb9517ed5d5ffaa377c60618a Mon Sep 17 00:00:00 2001
From: Kai Chen <chenkaidev@gmail.com>
Date: Wed, 17 Oct 2018 00:11:26 +0800
Subject: [PATCH] bug fix for compiling nms op
---
mmdet/ops/nms/Makefile | 2 +-
mmdet/ops/nms/setup.py | 33 ++++++++-------------------------
2 files changed, 9 insertions(+), 26 deletions(-)
diff --git a/mmdet/ops/nms/Makefile b/mmdet/ops/nms/Makefile
index 39556dd..af511f3 100644
--- a/mmdet/ops/nms/Makefile
+++ b/mmdet/ops/nms/Makefile
@@ -5,4 +5,4 @@ all:
$(PYTHON) setup.py build_ext --inplace
clean:
- rm *.so
+ rm -f *.so
diff --git a/mmdet/ops/nms/setup.py b/mmdet/ops/nms/setup.py
index 98bf57c..a8fe373 100644
--- a/mmdet/ops/nms/setup.py
+++ b/mmdet/ops/nms/setup.py
@@ -1,41 +1,24 @@
-import os
-from distutils.core import setup
-from distutils.extension import Extension
+import os.path as osp
+from distutils.core import setup, Extension
import numpy as np
from Cython.Build import cythonize
from Cython.Distutils import build_ext
-CUDA_ROOT = '/usr/local/cuda'
-CUDA = {
- "include": os.path.join(CUDA_ROOT, 'include'),
- "lib": os.path.join(CUDA_ROOT, 'lib64'),
- "nvcc": os.path.join(CUDA_ROOT, 'bin', "nvcc")
-}
-
-inc_dirs = [CUDA['include'], np.get_include()]
-
-lib_dirs = [CUDA['lib']]
-
# extensions
ext_args = dict(
- include_dirs=inc_dirs,
- library_dirs=lib_dirs,
+ include_dirs=[np.get_include()],
language='c++',
- libraries=['cudart'],
extra_compile_args={
- "cc": ['-Wno-unused-function', '-Wno-write-strings'],
- "nvcc": [
- '-arch=sm_52', '--ptxas-options=-v', '-c', '--compiler-options',
- '-fPIC'
- ],
+ 'cc': ['-Wno-unused-function', '-Wno-write-strings'],
+ 'nvcc': ['-c', '--compiler-options', '-fPIC'],
},
)
extensions = [
Extension('cpu_nms', ['cpu_nms.pyx'], **ext_args),
- Extension('gpu_nms', ['gpu_nms.pyx', 'nms_kernel.cu'], **ext_args),
Extension('cpu_soft_nms', ['cpu_soft_nms.pyx'], **ext_args),
+ Extension('gpu_nms', ['gpu_nms.pyx', 'nms_kernel.cu'], **ext_args),
]
@@ -59,9 +42,9 @@ def customize_compiler_for_nvcc(self):
# object but distutils doesn't have the ability to change compilers
# based on source extension: we add it.
def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
- if os.path.splitext(src)[1] == '.cu':
+ if osp.splitext(src)[1] == '.cu':
# use the cuda for .cu files
- self.set_executable('compiler_so', CUDA['nvcc'])
+ self.set_executable('compiler_so', 'nvcc')
# use only a subset of the extra_postargs, which are 1-1 translated
# from the extra_compile_args in the Extension class
postargs = extra_postargs['nvcc']
--
GitLab