Refactor compilation into single setup.py (#881)

* Implement all extension building in main setup.py

* Remove old build files

* Adapt installation guide

* Refactor to comply with flake8

* Move imports into functions & refactor module paths

* Format setup.py with yapf

* Move cython/numpy import back to head of file

* Add note to use verbose mode for pip in INSTALL.md

INSTALL.md

@@ -37,17 +37,11 @@ git clone https://github.com/open-mmlab/mmdetection.git
 cd mmdetection
 ```
 
-d. Compile cuda extensions.
-
-```shell
-./compile.sh
-```
-
-e. Install mmdetection (other dependencies will be installed automatically).
+d. Install mmdetection (other dependencies will be installed automatically).
 
 ```shell
 python setup.py develop
-# or "pip install -e ."
+# or "pip install -v -e ."
 ```
 
 Note:
@@ -83,7 +77,7 @@ mmdetection
 a script for setting up mmdetection with conda.
 
 ### Notice
-You can run `python(3) setup.py develop` or `pip install -e .` to install mmdetection if you want to make modifications to it frequently.
+You can run `python(3) setup.py develop` or `pip install -v -e .` to install mmdetection if you want to make modifications to it frequently.
 
 If there are more than one mmdetection on your machine, and you want to use them alternatively.
 Please insert the following code to the main file

compile.sh

-#!/usr/bin/env bash
-
-PYTHON=${PYTHON:-"python"}
-
-echo "Building roi align op..."
-cd mmdet/ops/roi_align
-if [ -d "build" ]; then
-    rm -r build
-fi
-$PYTHON setup.py build_ext --inplace
-
-echo "Building roi pool op..."
-cd ../roi_pool
-if [ -d "build" ]; then
-    rm -r build
-fi
-$PYTHON setup.py build_ext --inplace
-
-echo "Building nms op..."
-cd ../nms
-if [ -d "build" ]; then
-    rm -r build
-fi
-$PYTHON setup.py build_ext --inplace
-
-echo "Building dcn..."
-cd ../dcn
-if [ -d "build" ]; then
-    rm -r build
-fi
-$PYTHON setup.py build_ext --inplace
-
-echo "Building sigmoid focal loss op..."
-cd ../sigmoid_focal_loss
-if [ -d "build" ]; then
-    rm -r build
-fi
-$PYTHON setup.py build_ext --inplace
-
-echo "Building masked conv op..."
-cd ../masked_conv
-if [ -d "build" ]; then
-    rm -r build
-fi
-$PYTHON setup.py build_ext --inplace

mmdet/ops/dcn/setup.py

-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='deform_conv',
-    ext_modules=[
-        CUDAExtension('deform_conv_cuda', [
-            'src/deform_conv_cuda.cpp',
-            'src/deform_conv_cuda_kernel.cu',
-        ]),
-        CUDAExtension(
-            'deform_pool_cuda',
-            ['src/deform_pool_cuda.cpp', 'src/deform_pool_cuda_kernel.cu']),
-    ],
-    cmdclass={'build_ext': BuildExtension})

mmdet/ops/masked_conv/setup.py

-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='masked_conv2d_cuda',
-    ext_modules=[
-        CUDAExtension('masked_conv2d_cuda', [
-            'src/masked_conv2d_cuda.cpp',
-            'src/masked_conv2d_kernel.cu',
-        ]),
-    ],
-    cmdclass={'build_ext': BuildExtension})

mmdet/ops/nms/setup.py

-import os.path as osp
-from setuptools import setup, Extension
-
-import numpy as np
-from Cython.Build import cythonize
-from Cython.Distutils import build_ext
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-ext_args = dict(
-    include_dirs=[np.get_include()],
-    language='c++',
-    extra_compile_args={
-        'cc': ['-Wno-unused-function', '-Wno-write-strings'],
-        'nvcc': ['-c', '--compiler-options', '-fPIC'],
-    },
-)
-
-extensions = [
-    Extension('soft_nms_cpu', ['src/soft_nms_cpu.pyx'], **ext_args),
-]
-
-
-def customize_compiler_for_nvcc(self):
-    """inject deep into distutils to customize how the dispatch
-    to cc/nvcc works.
-    If you subclass UnixCCompiler, it's not trivial to get your subclass
-    injected in, and still have the right customizations (i.e.
-    distutils.sysconfig.customize_compiler) run on it. So instead of going
-    the OO route, I have this. Note, it's kindof like a wierd functional
-    subclassing going on."""
-
-    # tell the compiler it can processes .cu
-    self.src_extensions.append('.cu')
-
-    # save references to the default compiler_so and _comple methods
-    default_compiler_so = self.compiler_so
-    super = self._compile
-
-    # now redefine the _compile method. This gets executed for each
-    # 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 osp.splitext(src)[1] == '.cu':
-            # use the cuda for .cu files
-            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']
-        else:
-            postargs = extra_postargs['cc']
-
-        super(obj, src, ext, cc_args, postargs, pp_opts)
-        # reset the default compiler_so, which we might have changed for cuda
-        self.compiler_so = default_compiler_so
-
-    # inject our redefined _compile method into the class
-    self._compile = _compile
-
-
-class custom_build_ext(build_ext):
-
-    def build_extensions(self):
-        customize_compiler_for_nvcc(self.compiler)
-        build_ext.build_extensions(self)
-
-
-setup(
-    name='soft_nms',
-    cmdclass={'build_ext': custom_build_ext},
-    ext_modules=cythonize(extensions),
-)
-
-setup(
-    name='nms_cuda',
-    ext_modules=[
-        CUDAExtension('nms_cuda', [
-            'src/nms_cuda.cpp',
-            'src/nms_kernel.cu',
-        ]),
-        CUDAExtension('nms_cpu', [
-            'src/nms_cpu.cpp',
-        ]),
-    ],
-    cmdclass={'build_ext': BuildExtension})

mmdet/ops/roi_align/setup.py

-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='roi_align_cuda',
-    ext_modules=[
-        CUDAExtension('roi_align_cuda', [
-            'src/roi_align_cuda.cpp',
-            'src/roi_align_kernel.cu',
-        ]),
-    ],
-    cmdclass={'build_ext': BuildExtension})

mmdet/ops/roi_pool/setup.py

-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='roi_pool',
-    ext_modules=[
-        CUDAExtension('roi_pool_cuda', [
-            'src/roi_pool_cuda.cpp',
-            'src/roi_pool_kernel.cu',
-        ])
-    ],
-    cmdclass={'build_ext': BuildExtension})

mmdet/ops/sigmoid_focal_loss/setup.py

-from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-setup(
-    name='SigmoidFocalLoss',
-    ext_modules=[
-        CUDAExtension('sigmoid_focal_loss_cuda', [
-            'src/sigmoid_focal_loss.cpp',
-            'src/sigmoid_focal_loss_cuda.cu',
-        ]),
-    ],
-    cmdclass={'build_ext': BuildExtension})

setup.py

 import os
 import subprocess
 import time
-from setuptools import find_packages, setup
+from setuptools import Extension, find_packages, setup
+
+import numpy as np
+from Cython.Build import cythonize
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
 
 def readme():
@@ -80,6 +84,26 @@ def get_version():
     return locals()['__version__']
 
 
+def make_cuda_ext(name, module, sources):
+
+    return CUDAExtension(
+        name='{}.{}'.format(module, name),
+        sources=[os.path.join(*module.split('.'), p) for p in sources])
+
+
+def make_cython_ext(name, module, sources):
+    extension = Extension(
+        '{}.{}'.format(module, name),
+        [os.path.join(*module.split('.'), p) for p in sources],
+        include_dirs=[np.get_include()],
+        language='c++',
+        extra_compile_args={
+            'cxx': ['-Wno-unused-function', '-Wno-write-strings']
+        })
+    extension, = cythonize(extension)
+    return extension
+
+
 if __name__ == '__main__':
     write_version_py()
     setup(
@@ -103,10 +127,60 @@ if __name__ == '__main__':
             'Programming Language :: Python :: 3.6',
         ],
         license='Apache License 2.0',
-        setup_requires=['pytest-runner'],
+        setup_requires=['pytest-runner', 'cython', 'numpy'],
         tests_require=['pytest'],
         install_requires=[
             'mmcv>=0.2.6', 'numpy', 'matplotlib', 'six', 'terminaltables',
-            'pycocotools'
+            'pycocotools', 'torch>=1.1'
+        ],
+        ext_modules=[
+            make_cython_ext(
+                name='soft_nms_cpu',
+                module='mmdet.ops.nms',
+                sources=['src/soft_nms_cpu.pyx']),
+            make_cuda_ext(
+                name='roi_align_cuda',
+                module='mmdet.ops.roi_align',
+                sources=['src/roi_align_cuda.cpp', 'src/roi_align_kernel.cu']),
+            make_cuda_ext(
+                name='roi_pool_cuda',
+                module='mmdet.ops.roi_pool',
+                sources=['src/roi_pool_cuda.cpp', 'src/roi_pool_kernel.cu']),
+            make_cuda_ext(
+                name='nms_cpu',
+                module='mmdet.ops.nms',
+                sources=['src/nms_cpu.cpp']),
+            make_cuda_ext(
+                name='nms_cuda',
+                module='mmdet.ops.nms',
+                sources=['src/nms_cuda.cpp', 'src/nms_kernel.cu']),
+            make_cuda_ext(
+                name='deform_conv_cuda',
+                module='mmdet.ops.dcn',
+                sources=[
+                    'src/deform_conv_cuda.cpp',
+                    'src/deform_conv_cuda_kernel.cu'
+                ]),
+            make_cuda_ext(
+                name='deform_pool_cuda',
+                module='mmdet.ops.dcn',
+                sources=[
+                    'src/deform_pool_cuda.cpp',
+                    'src/deform_pool_cuda_kernel.cu'
+                ]),
+            make_cuda_ext(
+                name='sigmoid_focal_loss_cuda',
+                module='mmdet.ops.sigmoid_focal_loss',
+                sources=[
+                    'src/sigmoid_focal_loss.cpp',
+                    'src/sigmoid_focal_loss_cuda.cu'
+                ]),
+            make_cuda_ext(
+                name='masked_conv2d_cuda',
+                module='mmdet.ops.masked_conv',
+                sources=[
+                    'src/masked_conv2d_cuda.cpp', 'src/masked_conv2d_kernel.cu'
+                ]),
         ],
+        cmdclass={'build_ext': BuildExtension},
         zip_safe=False)