setup.py 15.4 KB
Newer Older
1
2
import glob
import os
3
import platform
Rinat Shigapov's avatar
Rinat Shigapov committed
4
import re
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'
37
    with open(version_file, 'r', encoding='utf-8') as f:
38
39
40
41
        exec(compile(f.read(), version_file, 'exec'))
    return locals()['__version__']


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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
        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):
        with open(fpath, 'r') as f:
            for line in f.readlines():
                line = line.strip()
                if line and not line.startswith('#'):
                    for info in parse_line(line):
                        yield info

    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 = []

137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
    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': []}

152
        include_path = os.path.abspath('./mmcv/ops/csrc/common/cuda')
153
154
155
156
157
158
159
160
161
162
163
        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 []
164
165
166
        # prevent cub/thrust conflict with other python library
        # More context See issues #1454
        extra_compile_args['nvcc'] += ['-Xcompiler=-fno-gnu-unique']
167
168
169
170
171
172
173
174
175
176
177
178
179
180
        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)

181
182
183
    if os.getenv('MMCV_WITH_OPS', '0') == '0':
        return extensions

Yuanhao Zhu's avatar
Yuanhao Zhu committed
184
185
186
    if EXT_TYPE == 'parrots':
        ext_name = 'mmcv._ext'
        from parrots.utils.build_extension import Extension
187
188
189
        # new parrots op impl do not use MMCV_USE_PARROTS
        # define_macros = [('MMCV_USE_PARROTS', None)]
        define_macros = []
190
191
        include_dirs = []
        op_files = glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cu') +\
192
            glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') +\
193
            glob.glob('./mmcv/ops/csrc/parrots/*.cpp')
194
195
        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
196
        cuda_args = os.getenv('MMCV_CUDA_ARGS')
197
        extra_compile_args = {
198
199
            'nvcc': [cuda_args, '-std=c++14'] if cuda_args else ['-std=c++14'],
            'cxx': ['-std=c++14'],
200
201
202
203
204
205
206
207
        }
        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
208
209
210
        ext_ops = Extension(
            name=ext_name,
            sources=op_files,
211
            include_dirs=include_dirs,
Yuanhao Zhu's avatar
Yuanhao Zhu committed
212
            define_macros=define_macros,
213
214
215
            extra_compile_args=extra_compile_args,
            cuda=True,
            pytorch=True)
Yuanhao Zhu's avatar
Yuanhao Zhu committed
216
217
        extensions.append(ext_ops)
    elif EXT_TYPE == 'pytorch':
218
        ext_name = 'mmcv._ext'
219
220
        from torch.utils.cpp_extension import CppExtension, CUDAExtension

Yuanhao Zhu's avatar
Yuanhao Zhu committed
221
        # prevent ninja from using too many resources
lizz's avatar
lizz committed
222
223
224
225
226
227
228
229
        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
230
        define_macros = []
231
232
233
234
235
236
237
238
239

        # 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': []}

240
241
242
243
244
245
246
247
248
249
        # 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 platform.system() != 'Windows':
            extra_compile_args['cxx'] = ['-std=c++14']

250
        include_dirs = []
Yuanhao Zhu's avatar
Yuanhao Zhu committed
251

252
        is_rocm_pytorch = False
253
        try:
254
255
256
            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
257
258
        except ImportError:
            pass
259

260
        project_dir = 'mmcv/ops/csrc/'
261
262
263
264
        if is_rocm_pytorch:
            from torch.utils.hipify import hipify_python

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

297
298
299
300
301
302
303
304
        # 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':
305
306
            extra_compile_args['nvcc'] += ['-std=c++14']

Yuanhao Zhu's avatar
Yuanhao Zhu committed
307
308
309
        ext_ops = extension(
            name=ext_name,
            sources=op_files,
310
            include_dirs=include_dirs,
Yuanhao Zhu's avatar
Yuanhao Zhu committed
311
312
313
            define_macros=define_macros,
            extra_compile_args=extra_compile_args)
        extensions.append(ext_ops)
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

    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/*')
tangyanf's avatar
tangyanf committed
339
            include_dirs += include_paths(cuda=True)
340
341
            library_dirs += library_paths(cuda=True)
        else:
tangyanf's avatar
tangyanf committed
342
            include_dirs += include_paths(cuda=False)
343
344
345
346
347
348
349
350
351
352
353
354
355
356
            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)

357
358
    return extensions

359

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