setup.py 15 KB
Newer Older
1
2
import glob
import os
3
import platform
Rinat Shigapov's avatar
Rinat Shigapov committed
4
import re
limm's avatar
limm committed
5
from pkg_resources import DistributionNotFound, get_distribution
6
from setuptools import find_packages, setup
Yuanhao Zhu's avatar
Yuanhao Zhu committed
7
8
9
10
11
12
13
14
15
16

EXT_TYPE = ''
try:
    import torch
    if torch.__version__ == 'parrots':
        from parrots.utils.build_extension import BuildExtension
        EXT_TYPE = 'parrots'
    else:
        from torch.utils.cpp_extension import BuildExtension
        EXT_TYPE = 'pytorch'
17
    cmd_class = {'build_ext': BuildExtension}
Yuanhao Zhu's avatar
Yuanhao Zhu committed
18
except ModuleNotFoundError:
19
    cmd_class = {}
Yuanhao Zhu's avatar
Yuanhao Zhu committed
20
    print('Skip building ext ops due to the absence of torch.')
21

Rinat Shigapov's avatar
Rinat Shigapov committed
22
23

def choose_requirement(primary, secondary):
Kai Chen's avatar
Kai Chen committed
24
25
    """If some version of primary requirement installed, return primary, else
    return secondary."""
Rinat Shigapov's avatar
Rinat Shigapov committed
26
27
28
29
30
31
32
33
34
    try:
        name = re.split(r'[!<>=]', primary)[0]
        get_distribution(name)
    except DistributionNotFound:
        return secondary

    return str(primary)


35
36
def get_version():
    version_file = 'mmcv/version.py'
limm's avatar
limm committed
37
    with open(version_file, 'r', encoding='utf-8') as f:
38
        exec(compile(f.read(), version_file, 'exec'))
limm's avatar
limm committed
39
    return locals()['__version__']
40
41


42
def parse_requirements(fname='requirements/runtime.txt', with_version=True):
Kai Chen's avatar
Kai Chen committed
43
    """Parse the package dependencies listed in a requirements file but strips
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
    specific versioning information.

    Args:
        fname (str): path to requirements file
        with_version (bool, default=False): if True include version specs

    Returns:
        List[str]: list of requirements items

    CommandLine:
        python -c "import setup; print(setup.parse_requirements())"
    """
    import sys
    from os.path import exists
    require_fpath = fname

    def parse_line(line):
Kai Chen's avatar
Kai Chen committed
61
        """Parse information from a line in a requirements text file."""
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
        if line.startswith('-r '):
            # Allow specifying requirements in other files
            target = line.split(' ')[1]
            for info in parse_require_file(target):
                yield info
        else:
            info = {'line': line}
            if line.startswith('-e '):
                info['package'] = line.split('#egg=')[1]
            else:
                # Remove versioning from the package
                pat = '(' + '|'.join(['>=', '==', '>']) + ')'
                parts = re.split(pat, line, maxsplit=1)
                parts = [p.strip() for p in parts]

                info['package'] = parts[0]
                if len(parts) > 1:
                    op, rest = parts[1:]
                    if ';' in rest:
                        # Handle platform specific dependencies
                        # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies
                        version, platform_deps = map(str.strip,
                                                     rest.split(';'))
                        info['platform_deps'] = platform_deps
                    else:
                        version = rest  # NOQA
                    info['version'] = (op, version)
            yield info

    def parse_require_file(fpath):
limm's avatar
limm committed
92
        with open(fpath, 'r') as f:
93
94
95
            for line in f.readlines():
                line = line.strip()
                if line and not line.startswith('#'):
limm's avatar
limm committed
96
97
                    for info in parse_line(line):
                        yield info
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

    def gen_packages_items():
        if exists(require_fpath):
            for info in parse_require_file(require_fpath):
                parts = [info['package']]
                if with_version and 'version' in info:
                    parts.extend(info['version'])
                if not sys.version.startswith('3.4'):
                    # apparently package_deps are broken in 3.4
                    platform_deps = info.get('platform_deps')
                    if platform_deps is not None:
                        parts.append(';' + platform_deps)
                item = ''.join(parts)
                yield item

    packages = list(gen_packages_items())
    return packages


install_requires = parse_requirements()
118
119
120
121
122
123
124
125
126
127
128
129
130
131

try:
    # OpenCV installed via conda.
    import cv2  # NOQA: F401
    major, minor, *rest = cv2.__version__.split('.')
    if int(major) < 3:
        raise RuntimeError(
            f'OpenCV >=3 is required but {cv2.__version__} is installed')
except ImportError:
    # If first not installed install second package
    CHOOSE_INSTALL_REQUIRES = [('opencv-python-headless>=3',
                                'opencv-python>=3')]
    for main, secondary in CHOOSE_INSTALL_REQUIRES:
        install_requires.append(choose_requirement(main, secondary))
132

Kai Chen's avatar
Kai Chen committed
133

134
135
136
def get_extensions():
    extensions = []

limm's avatar
limm committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
    if os.getenv('MMCV_WITH_TRT', '0') != '0':
        ext_name = 'mmcv._ext_trt'
        from torch.utils.cpp_extension import include_paths, library_paths
        library_dirs = []
        libraries = []
        include_dirs = []
        tensorrt_path = os.getenv('TENSORRT_DIR', '0')
        tensorrt_lib_path = glob.glob(
            os.path.join(tensorrt_path, 'targets', '*', 'lib'))[0]
        library_dirs += [tensorrt_lib_path]
        libraries += ['nvinfer', 'nvparsers', 'nvinfer_plugin']
        libraries += ['cudart']
        define_macros = []
        extra_compile_args = {'cxx': []}

        include_path = os.path.abspath('./mmcv/ops/csrc/common/cuda')
        include_trt_path = os.path.abspath('./mmcv/ops/csrc/tensorrt')
        include_dirs.append(include_path)
        include_dirs.append(include_trt_path)
        include_dirs.append(os.path.join(tensorrt_path, 'include'))
        include_dirs += include_paths(cuda=True)

        op_files = glob.glob('./mmcv/ops/csrc/tensorrt/plugins/*')
        define_macros += [('MMCV_WITH_CUDA', None)]
        define_macros += [('MMCV_WITH_TRT', None)]
        cuda_args = os.getenv('MMCV_CUDA_ARGS')
        extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
        library_dirs += library_paths(cuda=True)

        from setuptools import Extension
        ext_ops = Extension(
            name=ext_name,
            sources=op_files,
            include_dirs=include_dirs,
            define_macros=define_macros,
            extra_compile_args=extra_compile_args,
            language='c++',
            library_dirs=library_dirs,
            libraries=libraries)
        extensions.append(ext_ops)

    if os.getenv('MMCV_WITH_OPS', '0') == '0':
179
180
        return extensions

Yuanhao Zhu's avatar
Yuanhao Zhu committed
181
182
183
    if EXT_TYPE == 'parrots':
        ext_name = 'mmcv._ext'
        from parrots.utils.build_extension import Extension
184
185
186
        # new parrots op impl do not use MMCV_USE_PARROTS
        # define_macros = [('MMCV_USE_PARROTS', None)]
        define_macros = []
187
188
        include_dirs = []
        op_files = glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cu') +\
189
            glob.glob('./mmcv/ops/csrc/parrots/*.cpp')
190
191
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/cuda'))
Yuanhao Zhu's avatar
Yuanhao Zhu committed
192
        cuda_args = os.getenv('MMCV_CUDA_ARGS')
193
        extra_compile_args = {
limm's avatar
limm committed
194
195
            'nvcc': [cuda_args] if cuda_args else [],
            'cxx': [],
196
197
198
199
200
201
202
203
        }
        if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':
            define_macros += [('MMCV_WITH_CUDA', None)]
            extra_compile_args['nvcc'] += [
                '-D__CUDA_NO_HALF_OPERATORS__',
                '-D__CUDA_NO_HALF_CONVERSIONS__',
                '-D__CUDA_NO_HALF2_OPERATORS__',
            ]
Yuanhao Zhu's avatar
Yuanhao Zhu committed
204
205
206
        ext_ops = Extension(
            name=ext_name,
            sources=op_files,
207
            include_dirs=include_dirs,
Yuanhao Zhu's avatar
Yuanhao Zhu committed
208
            define_macros=define_macros,
209
210
211
            extra_compile_args=extra_compile_args,
            cuda=True,
            pytorch=True)
Yuanhao Zhu's avatar
Yuanhao Zhu committed
212
213
        extensions.append(ext_ops)
    elif EXT_TYPE == 'pytorch':
214
        ext_name = 'mmcv._ext'
215
        from torch.utils.cpp_extension import CppExtension, CUDAExtension
limm's avatar
limm committed
216
        from fastpt import CUDAExtension
Yuanhao Zhu's avatar
Yuanhao Zhu committed
217
        # prevent ninja from using too many resources
lizz's avatar
lizz committed
218
219
220
221
222
223
224
225
        try:
            import psutil
            num_cpu = len(psutil.Process().cpu_affinity())
            cpu_use = max(4, num_cpu - 1)
        except (ModuleNotFoundError, AttributeError):
            cpu_use = 4

        os.environ.setdefault('MAX_JOBS', str(cpu_use))
Yuanhao Zhu's avatar
Yuanhao Zhu committed
226
        define_macros = []
227
228
229
230
231
232
233
234
235

        # Before PyTorch1.8.0, when compiling CUDA code, `cxx` is a
        # required key passed to PyTorch. Even if there is no flag passed
        # to cxx, users also need to pass an empty list to PyTorch.
        # Since PyTorch1.8.0, it has a default value so users do not need
        # to pass an empty list anymore.
        # More details at https://github.com/pytorch/pytorch/pull/45956
        extra_compile_args = {'cxx': []}

limm's avatar
limm committed
236
237
238
239
240
241
242
        # Since the PR (https://github.com/open-mmlab/mmcv/pull/1463) uses
        # c++14 features, the argument ['std=c++14'] must be added here.
        # However, in the windows environment, some standard libraries
        # will depend on c++17 or higher. In fact, for the windows
        # environment, the compiler will choose the appropriate compiler
        # to compile those cpp files, so there is no need to add the
        # argument
243
244
245
        if platform.system() != 'Windows':
            extra_compile_args['cxx'] = ['-std=c++14']

246
        include_dirs = []
Yuanhao Zhu's avatar
Yuanhao Zhu committed
247

248
        is_rocm_pytorch = False
249
        try:
250
251
252
            from torch.utils.cpp_extension import ROCM_HOME
            is_rocm_pytorch = True if ((torch.version.hip is not None) and
                                       (ROCM_HOME is not None)) else False
253
254
        except ImportError:
            pass
255

limm's avatar
limm committed
256
257
258
259
260
261
262
263
264
265
266
267
        project_dir = 'mmcv/ops/csrc/'
        if is_rocm_pytorch and False:
            #from torch.utils.hipify import hipify_python
            from fastpt import hipify_python

            hipify_python.hipify(
                project_directory=project_dir,
                output_directory=project_dir,
                includes='mmcv/ops/csrc/*',
                show_detailed=True,
                is_pytorch_extension=True,
            )
Yuanhao Zhu's avatar
Yuanhao Zhu committed
268
            define_macros += [('MMCV_WITH_CUDA', None)]
limm's avatar
limm committed
269
270
271
272
273
274
275
276
277
278
279
            define_macros += [('HIP_DIFF', None)]
            cuda_args = os.getenv('MMCV_CUDA_ARGS')
            extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/hip/*') \
                + glob.glob('./mmcv/ops/csrc/pytorch/cpu/hip/*')
            extension = CUDAExtension
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/hip'))
        elif torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':
            define_macros += [('MMCV_WITH_CUDA', None)]
            if is_rocm_pytorch:
                define_macros += [('HIP_DIFF', None)
zhuyuanhao's avatar
zhuyuanhao committed
280
            cuda_args = os.getenv('MMCV_CUDA_ARGS')
Yuanhao Zhu's avatar
Yuanhao Zhu committed
281
            extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
282
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
283
284
285
                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cu') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cpp')
Yuanhao Zhu's avatar
Yuanhao Zhu committed
286
            extension = CUDAExtension
287
288
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/cuda'))
289
        else:
limm's avatar
limm committed
290
            print(f'Compiling {ext_name} without CUDA')
291
292
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp')
Yuanhao Zhu's avatar
Yuanhao Zhu committed
293
            extension = CppExtension
294
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
Yuanhao Zhu's avatar
Yuanhao Zhu committed
295

296
297
298
299
300
301
302
303
        # Since the PR (https://github.com/open-mmlab/mmcv/pull/1463) uses
        # c++14 features, the argument ['std=c++14'] must be added here.
        # However, in the windows environment, some standard libraries
        # will depend on c++17 or higher. In fact, for the windows
        # environment, the compiler will choose the appropriate compiler
        # to compile those cpp files, so there is no need to add the
        # argument
        if 'nvcc' in extra_compile_args and platform.system() != 'Windows':
304
305
            extra_compile_args['nvcc'] += ['-std=c++14']

Yuanhao Zhu's avatar
Yuanhao Zhu committed
306
307
308
        ext_ops = extension(
            name=ext_name,
            sources=op_files,
309
            include_dirs=include_dirs,
Yuanhao Zhu's avatar
Yuanhao Zhu committed
310
311
312
            define_macros=define_macros,
            extra_compile_args=extra_compile_args)
        extensions.append(ext_ops)
limm's avatar
limm committed
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355

    if EXT_TYPE == 'pytorch' and os.getenv('MMCV_WITH_ORT', '0') != '0':
        ext_name = 'mmcv._ext_ort'
        from torch.utils.cpp_extension import library_paths, include_paths
        import onnxruntime
        library_dirs = []
        libraries = []
        include_dirs = []
        ort_path = os.getenv('ONNXRUNTIME_DIR', '0')
        library_dirs += [os.path.join(ort_path, 'lib')]
        libraries.append('onnxruntime')
        define_macros = []
        extra_compile_args = {'cxx': []}

        include_path = os.path.abspath('./mmcv/ops/csrc/onnxruntime')
        include_dirs.append(include_path)
        include_dirs.append(os.path.join(ort_path, 'include'))

        op_files = glob.glob('./mmcv/ops/csrc/onnxruntime/cpu/*')
        if onnxruntime.get_device() == 'GPU' or os.getenv('FORCE_CUDA',
                                                          '0') == '1':
            define_macros += [('MMCV_WITH_CUDA', None)]
            cuda_args = os.getenv('MMCV_CUDA_ARGS')
            extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
            op_files += glob.glob('./mmcv/ops/csrc/onnxruntime/gpu/*')
            include_dirs += include_paths(cuda=True)
            library_dirs += library_paths(cuda=True)
        else:
            include_dirs += include_paths(cuda=False)
            library_dirs += library_paths(cuda=False)

        from setuptools import Extension
        ext_ops = Extension(
            name=ext_name,
            sources=op_files,
            include_dirs=include_dirs,
            define_macros=define_macros,
            extra_compile_args=extra_compile_args,
            language='c++',
            library_dirs=library_dirs,
            libraries=libraries)
        extensions.append(ext_ops)

356
357
    return extensions

358

359
setup(
limm's avatar
limm committed
360
    name='mmcv' if os.getenv('MMCV_WITH_OPS', '0') == '0' else 'mmcv-full',
361
    version=get_version(),
362
    description='OpenMMLab Computer Vision Foundation',
363
364
    keywords='computer vision',
    packages=find_packages(),
365
    include_package_data=True,
366
367
    classifiers=[
        'Development Status :: 4 - Beta',
Kai Chen's avatar
Kai Chen committed
368
        'License :: OSI Approved :: Apache Software License',
369
370
        'Operating System :: OS Independent',
        'Programming Language :: Python :: 3',
limm's avatar
limm committed
371
        'Programming Language :: Python :: 3.6',
372
        'Programming Language :: Python :: 3.7',
373
        'Programming Language :: Python :: 3.8',
374
        'Programming Language :: Python :: 3.9',
375
376
377
        'Topic :: Utilities',
    ],
    url='https://github.com/open-mmlab/mmcv',
Zaida Zhou's avatar
Zaida Zhou committed
378
    author='MMCV Contributors',
Kai Chen's avatar
Kai Chen committed
379
    author_email='openmmlab@gmail.com',
limm's avatar
limm committed
380
381
    setup_requires=['pytest-runner'],
    tests_require=['pytest'],
382
    install_requires=install_requires,
383
    ext_modules=get_extensions(),
384
    cmdclass=cmd_class,
385
    zip_safe=False)