GrokDet/setup.py at master · GrokCV/GrokDet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
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
118
import glob
import os
import platform
from packaging.version import parse as parse_version
from setuptools import find_packages, setup

EXT_TYPE = ''
cmd_class = {}

try:
    import torch
    from torch.utils.cpp_extension import BuildExtension

    EXT_TYPE = 'pytorch'
    cmd_class = {'build_ext': BuildExtension}
except ModuleNotFoundError:
    print('Skip building ext ops due to the absence of torch.')


def get_version():
    version_file = 'mmcv/version.py'
    namespace = {}
    with open(version_file, encoding='utf-8') as f:
        exec(compile(f.read(), version_file, 'exec'), namespace)
    return namespace['__version__']


def get_extensions():
    if os.getenv('MMCV_WITH_OPS', '1') == '0':
        return []

    if EXT_TYPE != 'pytorch':
        return []

    import torch
    from torch.utils.cpp_extension import CppExtension, CUDAExtension

    ext_name = 'mmcv._ext'
    define_macros = []
    extra_compile_args = {'cxx': []}

    if platform.system() != 'Windows':
        if parse_version(torch.__version__) <= parse_version('1.12.1'):
            extra_compile_args['cxx'] = ['-std=c++14']
        else:
            extra_compile_args['cxx'] = ['-std=c++17']
    else:
        if parse_version(torch.__version__) <= parse_version('1.12.1'):
            extra_compile_args['cxx'] = ['/std:c++14']
        else:
            extra_compile_args['cxx'] = ['/std:c++17']

    include_dirs = []
    extra_objects = []
    library_dirs = []
    libraries = []
    extra_link_args = []

    if torch.cuda.is_available() 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/pytorch/*.cpp') + \
            glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
            glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cu') + \
            glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cpp')

        extension = CUDAExtension
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/pytorch'))
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/cuda'))

        if platform.system() != 'Windows':
            if parse_version(torch.__version__) <= parse_version('1.12.1'):
                extra_compile_args['nvcc'] += ['-std=c++14']
            else:
                extra_compile_args['nvcc'] += ['-std=c++17']
    else:
        print(f'Compiling {ext_name} only with CPU')
        op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
            glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp')
        extension = CppExtension
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))

    if not op_files:
        raise RuntimeError('No C++/CUDA source files found under mmcv/ops/csrc/pytorch')

    os.environ.setdefault('MAX_JOBS', os.getenv('MAX_JOBS', '8'))

    ext_ops = extension(
        name=ext_name,
        sources=op_files,
        include_dirs=include_dirs,
        define_macros=define_macros,
        extra_objects=extra_objects,
        extra_compile_args=extra_compile_args,
        library_dirs=library_dirs,
        libraries=libraries,
        extra_link_args=extra_link_args)

    return [ext_ops]


setup(
    name='mmcv',
    version=get_version(),
    description='Vendored MMCV ops for GrokDet-developing',
    keywords='computer vision',
    packages=find_packages(include=['mmcv', 'mmcv.*']),
    include_package_data=True,
    license_files=['LICENSE', 'NOTICE', 'THIRD_PARTY_LICENSES.md'],
    python_requires='>=3.8',
    install_requires=['packaging'],
    ext_modules=get_extensions(),
    cmdclass=cmd_class,
    zip_safe=False,
)