From 713e98bc1bfb842760df6be73b9520ee775d3c06 Mon Sep 17 00:00:00 2001
From: Daniel Suess <daniel@dsuess.me>
Date: Thu, 11 Jul 2019 18:11:56 +1000
Subject: [PATCH] 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                            | 12 +---
 compile.sh                            | 45 --------------
 mmdet/ops/dcn/setup.py                | 15 -----
 mmdet/ops/masked_conv/setup.py        | 12 ----
 mmdet/ops/nms/setup.py                | 84 ---------------------------
 mmdet/ops/roi_align/setup.py          | 12 ----
 mmdet/ops/roi_pool/setup.py           | 12 ----
 mmdet/ops/sigmoid_focal_loss/setup.py | 12 ----
 setup.py                              | 80 ++++++++++++++++++++++++-
 9 files changed, 80 insertions(+), 204 deletions(-)
 delete mode 100755 compile.sh
 delete mode 100644 mmdet/ops/dcn/setup.py
 delete mode 100644 mmdet/ops/masked_conv/setup.py
 delete mode 100644 mmdet/ops/nms/setup.py
 delete mode 100644 mmdet/ops/roi_align/setup.py
 delete mode 100644 mmdet/ops/roi_pool/setup.py
 delete mode 100644 mmdet/ops/sigmoid_focal_loss/setup.py

diff --git a/INSTALL.md b/INSTALL.md
index ef52f8f2..e48b1312 100644
--- a/INSTALL.md
+++ b/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
diff --git a/compile.sh b/compile.sh
deleted file mode 100755
index c3853f1e..00000000
--- a/compile.sh
+++ /dev/null
@@ -1,45 +0,0 @@
-#!/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
diff --git a/mmdet/ops/dcn/setup.py b/mmdet/ops/dcn/setup.py
deleted file mode 100644
index 96380181..00000000
--- a/mmdet/ops/dcn/setup.py
+++ /dev/null
@@ -1,15 +0,0 @@
-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})
diff --git a/mmdet/ops/masked_conv/setup.py b/mmdet/ops/masked_conv/setup.py
deleted file mode 100644
index fdff5f20..00000000
--- a/mmdet/ops/masked_conv/setup.py
+++ /dev/null
@@ -1,12 +0,0 @@
-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})
diff --git a/mmdet/ops/nms/setup.py b/mmdet/ops/nms/setup.py
deleted file mode 100644
index 28f3b4e9..00000000
--- a/mmdet/ops/nms/setup.py
+++ /dev/null
@@ -1,84 +0,0 @@
-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})
diff --git a/mmdet/ops/roi_align/setup.py b/mmdet/ops/roi_align/setup.py
deleted file mode 100644
index f02a5ea3..00000000
--- a/mmdet/ops/roi_align/setup.py
+++ /dev/null
@@ -1,12 +0,0 @@
-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})
diff --git a/mmdet/ops/roi_pool/setup.py b/mmdet/ops/roi_pool/setup.py
deleted file mode 100644
index 16991b88..00000000
--- a/mmdet/ops/roi_pool/setup.py
+++ /dev/null
@@ -1,12 +0,0 @@
-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})
diff --git a/mmdet/ops/sigmoid_focal_loss/setup.py b/mmdet/ops/sigmoid_focal_loss/setup.py
deleted file mode 100644
index a70c6545..00000000
--- a/mmdet/ops/sigmoid_focal_loss/setup.py
+++ /dev/null
@@ -1,12 +0,0 @@
-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})
diff --git a/setup.py b/setup.py
index 6a875207..e40909a1 100644
--- a/setup.py
+++ b/setup.py
@@ -1,7 +1,11 @@
 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)
-- 
GitLab