setup.py 18 KB
Newer Older
1
import distutils.command.clean
2
import distutils.spawn
3
import glob
4
import os
5
import shutil
6
7
import subprocess
import sys
8
9

import torch
10
11
from pkg_resources import parse_version, get_distribution, DistributionNotFound
from setuptools import setup, find_packages
12
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME
soumith's avatar
soumith committed
13
14


Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
15
def read(*names, **kwargs):
16
    with open(os.path.join(os.path.dirname(__file__), *names), encoding=kwargs.get("encoding", "utf8")) as fp:
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
17
18
        return fp.read()

Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
19

20
21
22
23
24
25
26
def get_dist(pkgname):
    try:
        return get_distribution(pkgname)
    except DistributionNotFound:
        return None


27
28
cwd = os.path.dirname(os.path.abspath(__file__))

29
version_txt = os.path.join(cwd, "version.txt")
30
with open(version_txt) as f:
31
    version = f.readline().strip()
32
33
sha = "Unknown"
package_name = "torchvision"
Sasank Chilamkurthy's avatar
Sasank Chilamkurthy committed
34

35
try:
36
    sha = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=cwd).decode("ascii").strip()
37
38
39
except Exception:
    pass

40
41
42
43
if os.getenv("BUILD_VERSION"):
    version = os.getenv("BUILD_VERSION")
elif sha != "Unknown":
    version += "+" + sha[:7]
44
45
46


def write_version_file():
47
48
    version_path = os.path.join(cwd, "torchvision", "version.py")
    with open(version_path, "w") as f:
49
50
        f.write(f"__version__ = '{version}'\n")
        f.write(f"git_version = {repr(sha)}\n")
51
52
53
        f.write("from torchvision.extension import _check_cuda_version\n")
        f.write("if _check_cuda_version() > 0:\n")
        f.write("    cuda = _check_cuda_version()\n")
54
55


56
57
58
pytorch_dep = "torch"
if os.getenv("PYTORCH_VERSION"):
    pytorch_dep += "==" + os.getenv("PYTORCH_VERSION")
soumith's avatar
soumith committed
59

60
requirements = [
61
    "numpy",
62
    pytorch_dep,
63
64
]

65
66
# Excluding 8.3.* because of https://github.com/pytorch/vision/issues/4934
pillow_ver = " >= 5.3.0, !=8.3.*"
67
pillow_req = "pillow-simd" if get_dist("pillow-simd") is not None else "pillow"
68
69
requirements.append(pillow_req + pillow_ver)

70

71
72
def find_library(name, vision_include):
    this_dir = os.path.dirname(os.path.abspath(__file__))
73
    build_prefix = os.environ.get("BUILD_PREFIX", None)
74
75
76
77
78
79
    is_conda_build = build_prefix is not None

    library_found = False
    conda_installed = False
    lib_folder = None
    include_folder = None
80
    library_header = f"{name}.h"
81

82
    # Lookup in TORCHVISION_INCLUDE or in the package file
83
    package_path = [os.path.join(this_dir, "torchvision")]
84
85
86
87
88
89
90
    for folder in vision_include + package_path:
        candidate_path = os.path.join(folder, library_header)
        library_found = os.path.exists(candidate_path)
        if library_found:
            break

    if not library_found:
91
        print(f"Running build on conda-build: {is_conda_build}")
92
93
        if is_conda_build:
            # Add conda headers/libraries
94
95
96
97
98
            if os.name == "nt":
                build_prefix = os.path.join(build_prefix, "Library")
            include_folder = os.path.join(build_prefix, "include")
            lib_folder = os.path.join(build_prefix, "lib")
            library_header_path = os.path.join(include_folder, library_header)
99
100
101
            library_found = os.path.isfile(library_header_path)
            conda_installed = library_found
        else:
102
            # Check if using Anaconda to produce wheels
103
            conda = distutils.spawn.find_executable("conda")
104
            is_conda = conda is not None
105
            print(f"Running build on conda: {is_conda}")
106
107
108
            if is_conda:
                python_executable = sys.executable
                py_folder = os.path.dirname(python_executable)
109
110
                if os.name == "nt":
                    env_path = os.path.join(py_folder, "Library")
111
112
                else:
                    env_path = os.path.dirname(py_folder)
113
114
115
                lib_folder = os.path.join(env_path, "lib")
                include_folder = os.path.join(env_path, "include")
                library_header_path = os.path.join(include_folder, library_header)
116
117
118
119
                library_found = os.path.isfile(library_header_path)
                conda_installed = library_found

        if not library_found:
120
            if sys.platform == "linux":
121
122
                library_found = os.path.exists(f"/usr/include/{library_header}")
                library_found = library_found or os.path.exists(f"/usr/local/include/{library_header}")
123
124
125
126

    return library_found, conda_installed, include_folder, lib_folder


127
128
def get_extensions():
    this_dir = os.path.dirname(os.path.abspath(__file__))
129
    extensions_dir = os.path.join(this_dir, "torchvision", "csrc")
130

131
132
133
    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) + glob.glob(
        os.path.join(extensions_dir, "ops", "*.cpp")
    )
134
    source_cpu = (
135
136
137
        glob.glob(os.path.join(extensions_dir, "ops", "autograd", "*.cpp"))
        + glob.glob(os.path.join(extensions_dir, "ops", "cpu", "*.cpp"))
        + glob.glob(os.path.join(extensions_dir, "ops", "quantized", "cpu", "*.cpp"))
138
    )
139
140

    is_rocm_pytorch = False
141
142

    if torch.__version__ >= "1.5":
143
        from torch.utils.cpp_extension import ROCM_HOME
144

145
        is_rocm_pytorch = (torch.version.hip is not None) and (ROCM_HOME is not None)
146
147

    if is_rocm_pytorch:
148
        from torch.utils.hipify import hipify_python
149

150
151
152
        hipify_python.hipify(
            project_directory=this_dir,
            output_directory=this_dir,
153
            includes="torchvision/csrc/ops/cuda/*",
154
155
            show_detailed=True,
            is_pytorch_extension=True,
156
        )
157
        source_cuda = glob.glob(os.path.join(extensions_dir, "ops", "hip", "*.hip"))
158
        # Copy over additional files
159
160
        for file in glob.glob(r"torchvision/csrc/ops/cuda/*.h"):
            shutil.copy(file, "torchvision/csrc/ops/hip")
161
    else:
162
        source_cuda = glob.glob(os.path.join(extensions_dir, "ops", "cuda", "*.cu"))
163

164
    source_cuda += glob.glob(os.path.join(extensions_dir, "ops", "autocast", "*.cpp"))
165
166
167
168

    sources = main_file + source_cpu
    extension = CppExtension

169
    compile_cpp_tests = os.getenv("WITH_CPP_MODELS_TEST", "0") == "1"
170
    if compile_cpp_tests:
171
172
173
174
        test_dir = os.path.join(this_dir, "test")
        models_dir = os.path.join(this_dir, "torchvision", "csrc", "models")
        test_file = glob.glob(os.path.join(test_dir, "*.cpp"))
        source_models = glob.glob(os.path.join(models_dir, "*.cpp"))
175
176
177
178
179

        test_file = [os.path.join(test_dir, s) for s in test_file]
        source_models = [os.path.join(models_dir, s) for s in source_models]
        tests = test_file + source_models
        tests_include_dirs = [test_dir, models_dir]
Shahriar's avatar
Shahriar committed
180

181
182
    define_macros = []

183
184
185
186
    extra_compile_args = {"cxx": []}
    if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) or os.getenv(
        "FORCE_CUDA", "0"
    ) == "1":
187
188
        extension = CUDAExtension
        sources += source_cuda
189
        if not is_rocm_pytorch:
190
191
192
            define_macros += [("WITH_CUDA", None)]
            nvcc_flags = os.getenv("NVCC_FLAGS", "")
            if nvcc_flags == "":
193
194
                nvcc_flags = []
            else:
195
                nvcc_flags = nvcc_flags.split(" ")
Soumith Chintala's avatar
Soumith Chintala committed
196
        else:
197
            define_macros += [("WITH_HIP", None)]
198
            nvcc_flags = []
199
        extra_compile_args["nvcc"] = nvcc_flags
200

201
202
    if sys.platform == "win32":
        define_macros += [("torchvision_EXPORTS", None)]
203

204
        extra_compile_args["cxx"].append("/MP")
Francisco Massa's avatar
Francisco Massa committed
205

206
    debug_mode = os.getenv("DEBUG", "0") == "1"
207
208
    if debug_mode:
        print("Compile in debug mode")
209
210
        extra_compile_args["cxx"].append("-g")
        extra_compile_args["cxx"].append("-O0")
211
212
213
        if "nvcc" in extra_compile_args:
            # we have to remove "-OX" and "-g" flag if exists and append
            nvcc_flags = extra_compile_args["nvcc"]
214
            extra_compile_args["nvcc"] = [f for f in nvcc_flags if not ("-O" in f or "-g" in f)]
215
216
217
            extra_compile_args["nvcc"].append("-O0")
            extra_compile_args["nvcc"].append("-g")

218
219
    sources = [os.path.join(extensions_dir, s) for s in sources]

220
    include_dirs = [extensions_dir]
221
222
223

    ext_modules = [
        extension(
224
            "torchvision._C",
225
            sorted(sources),
226
227
            include_dirs=include_dirs,
            define_macros=define_macros,
Soumith Chintala's avatar
Soumith Chintala committed
228
            extra_compile_args=extra_compile_args,
229
        )
230
    ]
231
232
233
    if compile_cpp_tests:
        ext_modules.append(
            extension(
234
                "torchvision._C_tests",
235
236
237
238
239
240
                tests,
                include_dirs=tests_include_dirs,
                define_macros=define_macros,
                extra_compile_args=extra_compile_args,
            )
        )
241

242
    # ------------------- Torchvision extra extensions ------------------------
243
244
245
246
    vision_include = os.environ.get("TORCHVISION_INCLUDE", None)
    vision_library = os.environ.get("TORCHVISION_LIBRARY", None)
    vision_include = vision_include.split(os.pathsep) if vision_include is not None else []
    vision_library = vision_library.split(os.pathsep) if vision_library is not None else []
247
248
249
250
251
252
253
254
255
256
    include_dirs += vision_include
    library_dirs = vision_library

    # Image reading extension
    image_macros = []
    image_include = [extensions_dir]
    image_library = []
    image_link_flags = []

    # Locating libPNG
257
258
    libpng = distutils.spawn.find_executable("libpng-config")
    pngfix = distutils.spawn.find_executable("pngfix")
259
    png_found = libpng is not None or pngfix is not None
260
    print(f"PNG found: {png_found}")
261
262
263
    if png_found:
        if libpng is not None:
            # Linux / Mac
264
265
            png_version = subprocess.run([libpng, "--version"], stdout=subprocess.PIPE)
            png_version = png_version.stdout.strip().decode("utf-8")
266
            print(f"libpng version: {png_version}")
267
268
            png_version = parse_version(png_version)
            if png_version >= parse_version("1.6.0"):
269
270
271
272
                print("Building torchvision with PNG image support")
                png_lib = subprocess.run([libpng, "--libdir"], stdout=subprocess.PIPE)
                png_lib = png_lib.stdout.strip().decode("utf-8")
                if "disabled" not in png_lib:
273
                    image_library += [png_lib]
274
275
276
                png_include = subprocess.run([libpng, "--I_opts"], stdout=subprocess.PIPE)
                png_include = png_include.stdout.strip().decode("utf-8")
                _, png_include = png_include.split("-I")
277
                print(f"libpng include path: {png_include}")
278
                image_include += [png_include]
279
                image_link_flags.append("png")
280
            else:
281
                print("libpng installed version is less than 1.6.0, disabling PNG support")
282
283
284
                png_found = False
        else:
            # Windows
285
286
            png_lib = os.path.join(os.path.dirname(os.path.dirname(pngfix)), "lib")
            png_include = os.path.join(os.path.dirname(os.path.dirname(pngfix)), "include", "libpng16")
287
288
            image_library += [png_lib]
            image_include += [png_include]
289
            image_link_flags.append("libpng")
290

291
    # Locating libjpeg
292
    (jpeg_found, jpeg_conda, jpeg_include, jpeg_lib) = find_library("jpeglib", vision_include)
293

294
    print(f"JPEG found: {jpeg_found}")
295
296
    image_macros += [("PNG_FOUND", str(int(png_found)))]
    image_macros += [("JPEG_FOUND", str(int(jpeg_found)))]
297
    if jpeg_found:
298
299
        print("Building torchvision with JPEG image support")
        image_link_flags.append("jpeg")
300
301
302
303
        if jpeg_conda:
            image_library += [jpeg_lib]
            image_include += [jpeg_include]

304
305
306
    # Locating nvjpeg
    # Should be included in CUDA_HOME for CUDA >= 10.1, which is the minimum version we have in the CI
    nvjpeg_found = (
307
308
309
        extension is CUDAExtension
        and CUDA_HOME is not None
        and os.path.exists(os.path.join(CUDA_HOME, "include", "nvjpeg.h"))
310
311
    )

312
    print(f"NVJPEG found: {nvjpeg_found}")
313
    image_macros += [("NVJPEG_FOUND", str(int(nvjpeg_found)))]
314
    if nvjpeg_found:
315
316
317
318
319
320
321
322
323
        print("Building torchvision with NVJPEG image support")
        image_link_flags.append("nvjpeg")

    image_path = os.path.join(extensions_dir, "io", "image")
    image_src = (
        glob.glob(os.path.join(image_path, "*.cpp"))
        + glob.glob(os.path.join(image_path, "cpu", "*.cpp"))
        + glob.glob(os.path.join(image_path, "cuda", "*.cpp"))
    )
324

325
    if png_found or jpeg_found:
326
327
328
329
330
331
332
333
334
335
336
337
338
        ext_modules.append(
            extension(
                "torchvision.image",
                image_src,
                include_dirs=image_include + include_dirs + [image_path],
                library_dirs=image_library + library_dirs,
                define_macros=image_macros,
                libraries=image_link_flags,
                extra_compile_args=extra_compile_args,
            )
        )

    ffmpeg_exe = distutils.spawn.find_executable("ffmpeg")
339
    has_ffmpeg = ffmpeg_exe is not None
340
341
342
343
    # FIXME: Building torchvision with ffmpeg on MacOS or with Python 3.9
    # FIXME: causes crash. See the following GitHub issues for more details.
    # FIXME: https://github.com/pytorch/pytorch/issues/65000
    # FIXME: https://github.com/pytorch/vision/issues/3367
344
    if sys.platform != "linux" or (sys.version_info.major == 3 and sys.version_info.minor == 9):
345
        has_ffmpeg = False
346
347
    if has_ffmpeg:
        try:
348
349
350
            # This is to check if ffmpeg is installed properly.
            subprocess.check_output(["ffmpeg", "-version"])
        except subprocess.CalledProcessError:
351
            print("Error fetching ffmpeg version, ignoring ffmpeg.")
352
353
            has_ffmpeg = False

354
    print(f"FFmpeg found: {has_ffmpeg}")
355

356
    if has_ffmpeg:
357
        ffmpeg_libraries = {"libavcodec", "libavformat", "libavutil", "libswresample", "libswscale"}
358

359
360
        ffmpeg_bin = os.path.dirname(ffmpeg_exe)
        ffmpeg_root = os.path.dirname(ffmpeg_bin)
361
362
        ffmpeg_include_dir = os.path.join(ffmpeg_root, "include")
        ffmpeg_library_dir = os.path.join(ffmpeg_root, "lib")
363

364
365
366
        gcc = distutils.spawn.find_executable("gcc")
        platform_tag = subprocess.run([gcc, "-print-multiarch"], stdout=subprocess.PIPE)
        platform_tag = platform_tag.stdout.strip().decode("utf-8")
367
368
369

        if platform_tag:
            # Most probably a Debian-based distribution
370
371
            ffmpeg_include_dir = [ffmpeg_include_dir, os.path.join(ffmpeg_include_dir, platform_tag)]
            ffmpeg_library_dir = [ffmpeg_library_dir, os.path.join(ffmpeg_library_dir, platform_tag)]
372
373
374
375
376
377
378
379
        else:
            ffmpeg_include_dir = [ffmpeg_include_dir]
            ffmpeg_library_dir = [ffmpeg_library_dir]

        has_ffmpeg = True
        for library in ffmpeg_libraries:
            library_found = False
            for search_path in ffmpeg_include_dir + include_dirs:
380
                full_path = os.path.join(search_path, library, "*.h")
381
382
383
                library_found |= len(glob.glob(full_path)) > 0

            if not library_found:
384
                print(f"{library} header files were not found, disabling ffmpeg support")
385
386
387
                has_ffmpeg = False

    if has_ffmpeg:
388
389
        print(f"ffmpeg include path: {ffmpeg_include_dir}")
        print(f"ffmpeg library_dir: {ffmpeg_library_dir}")
390
391

        # TorchVision base decoder + video reader
392
        video_reader_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "video_reader")
393
        video_reader_src = glob.glob(os.path.join(video_reader_src_dir, "*.cpp"))
394
395
        base_decoder_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "decoder")
        base_decoder_src = glob.glob(os.path.join(base_decoder_src_dir, "*.cpp"))
396
        # Torchvision video API
397
        videoapi_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "video")
398
        videoapi_src = glob.glob(os.path.join(videoapi_src_dir, "*.cpp"))
399
        # exclude tests
400
        base_decoder_src = [x for x in base_decoder_src if "_test.cpp" not in x]
401

402
        combined_src = video_reader_src + base_decoder_src + videoapi_src
403

404
405
        ext_modules.append(
            CppExtension(
406
                "torchvision.video_reader",
407
                combined_src,
408
                include_dirs=[
409
                    base_decoder_src_dir,
410
                    video_reader_src_dir,
411
                    videoapi_src_dir,
412
                    extensions_dir,
413
                    *ffmpeg_include_dir,
414
                    *include_dirs,
415
                ],
416
                library_dirs=ffmpeg_library_dir + library_dirs,
417
                libraries=[
418
419
420
421
422
                    "avcodec",
                    "avformat",
                    "avutil",
                    "swresample",
                    "swscale",
423
                ],
424
425
                extra_compile_args=["-std=c++14"] if os.name != "nt" else ["/std:c++14", "/MP"],
                extra_link_args=["-std=c++14" if os.name != "nt" else "/std:c++14"],
426
427
            )
        )
428
429
430
431
432
433

    return ext_modules


class clean(distutils.command.clean.clean):
    def run(self):
434
        with open(".gitignore") as f:
435
            ignores = f.read()
436
            for wildcard in filter(None, ignores.split("\n")):
437
438
439
440
441
442
443
444
445
446
                for filename in glob.glob(wildcard):
                    try:
                        os.remove(filename)
                    except OSError:
                        shutil.rmtree(filename, ignore_errors=True)

        # It's an old-style class in Python 2.7...
        distutils.command.clean.clean.run(self)


447
if __name__ == "__main__":
448
    print(f"Building wheel {package_name}-{version}")
449
450
451

    write_version_file()

452
    with open("README.rst") as f:
453
454
455
456
457
458
        readme = f.read()

    setup(
        # Metadata
        name=package_name,
        version=version,
459
460
461
462
        author="PyTorch Core Team",
        author_email="soumith@pytorch.org",
        url="https://github.com/pytorch/vision",
        description="image and video datasets and models for torch deep learning",
463
        long_description=readme,
464
        license="BSD",
465
        # Package info
466
        packages=find_packages(exclude=("test",)),
467
        package_data={package_name: ["*.dll", "*.dylib", "*.so", "prototype/datasets/_builtin/*.categories"]},
468
469
470
471
472
473
        zip_safe=False,
        install_requires=requirements,
        extras_require={
            "scipy": ["scipy"],
        },
        ext_modules=get_extensions(),
474
        python_requires=">=3.6",
475
        cmdclass={
476
477
478
            "build_ext": BuildExtension.with_options(no_python_abi_suffix=True),
            "clean": clean,
        },
479
    )