setup.py 21.1 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
12
from pkg_resources import DistributionNotFound, get_distribution, parse_version
from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDA_HOME, CUDAExtension
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
    "typing_extensions",
62
    "numpy",
63
    "requests",
64
    pytorch_dep,
65
66
]

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

72

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

    library_found = False
    conda_installed = False
    lib_folder = None
    include_folder = None
82
    library_header = f"{name}.h"
83

84
    # Lookup in TORCHVISION_INCLUDE or in the package file
85
    package_path = [os.path.join(this_dir, "torchvision")]
86
87
88
89
90
91
92
    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:
93
        print(f"Running build on conda-build: {is_conda_build}")
94
95
        if is_conda_build:
            # Add conda headers/libraries
96
97
98
99
100
            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)
101
102
103
            library_found = os.path.isfile(library_header_path)
            conda_installed = library_found
        else:
104
            # Check if using Anaconda to produce wheels
105
            conda = shutil.which("conda")
106
            is_conda = conda is not None
107
            print(f"Running build on conda: {is_conda}")
108
109
110
            if is_conda:
                python_executable = sys.executable
                py_folder = os.path.dirname(python_executable)
111
112
                if os.name == "nt":
                    env_path = os.path.join(py_folder, "Library")
113
114
                else:
                    env_path = os.path.dirname(py_folder)
115
116
117
                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)
118
119
120
121
                library_found = os.path.isfile(library_header_path)
                conda_installed = library_found

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

    return library_found, conda_installed, include_folder, lib_folder


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

133
134
135
    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp")) + glob.glob(
        os.path.join(extensions_dir, "ops", "*.cpp")
    )
136
    source_cpu = (
137
138
139
        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"))
140
    )
141

142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
    print("Compiling extensions with following flags:")
    force_cuda = os.getenv("FORCE_CUDA", "0") == "1"
    print(f"  FORCE_CUDA: {force_cuda}")
    debug_mode = os.getenv("DEBUG", "0") == "1"
    print(f"  DEBUG: {debug_mode}")
    use_png = os.getenv("TORCHVISION_USE_PNG", "1") == "1"
    print(f"  TORCHVISION_USE_PNG: {use_png}")
    use_jpeg = os.getenv("TORCHVISION_USE_JPEG", "1") == "1"
    print(f"  TORCHVISION_USE_JPEG: {use_jpeg}")
    use_nvjpeg = os.getenv("TORCHVISION_USE_NVJPEG", "1") == "1"
    print(f"  TORCHVISION_USE_NVJPEG: {use_nvjpeg}")
    use_ffmpeg = os.getenv("TORCHVISION_USE_FFMPEG", "1") == "1"
    print(f"  TORCHVISION_USE_FFMPEG: {use_ffmpeg}")
    use_video_codec = os.getenv("TORCHVISION_USE_VIDEO_CODEC", "1") == "1"
    print(f"  TORCHVISION_USE_VIDEO_CODEC: {use_video_codec}")

    nvcc_flags = os.getenv("NVCC_FLAGS", "")
    print(f"  NVCC_FLAGS: {nvcc_flags}")

161
    is_rocm_pytorch = False
162
163

    if torch.__version__ >= "1.5":
164
        from torch.utils.cpp_extension import ROCM_HOME
165

166
        is_rocm_pytorch = (torch.version.hip is not None) and (ROCM_HOME is not None)
167
168

    if is_rocm_pytorch:
169
        from torch.utils.hipify import hipify_python
170

171
172
173
        hipify_python.hipify(
            project_directory=this_dir,
            output_directory=this_dir,
174
            includes="torchvision/csrc/ops/cuda/*",
175
176
            show_detailed=True,
            is_pytorch_extension=True,
177
        )
178
        source_cuda = glob.glob(os.path.join(extensions_dir, "ops", "hip", "*.hip"))
179
        # Copy over additional files
180
181
        for file in glob.glob(r"torchvision/csrc/ops/cuda/*.h"):
            shutil.copy(file, "torchvision/csrc/ops/hip")
182
    else:
183
        source_cuda = glob.glob(os.path.join(extensions_dir, "ops", "cuda", "*.cu"))
184

185
    source_cuda += glob.glob(os.path.join(extensions_dir, "ops", "autocast", "*.cpp"))
186
187
188
189
190
191

    sources = main_file + source_cpu
    extension = CppExtension

    define_macros = []

192
    extra_compile_args = {"cxx": []}
193
    if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) or force_cuda:
194
195
        extension = CUDAExtension
        sources += source_cuda
196
        if not is_rocm_pytorch:
197
198
            define_macros += [("WITH_CUDA", None)]
            if nvcc_flags == "":
199
200
                nvcc_flags = []
            else:
201
                nvcc_flags = nvcc_flags.split(" ")
Soumith Chintala's avatar
Soumith Chintala committed
202
        else:
203
            define_macros += [("WITH_HIP", None)]
204
            nvcc_flags = []
205
        extra_compile_args["nvcc"] = nvcc_flags
206

207
208
    if sys.platform == "win32":
        define_macros += [("torchvision_EXPORTS", None)]
209
        define_macros += [("USE_PYTHON", None)]
210
        extra_compile_args["cxx"].append("/MP")
Francisco Massa's avatar
Francisco Massa committed
211

212
    if debug_mode:
213
        print("Compiling in debug mode")
214
215
        extra_compile_args["cxx"].append("-g")
        extra_compile_args["cxx"].append("-O0")
216
217
218
        if "nvcc" in extra_compile_args:
            # we have to remove "-OX" and "-g" flag if exists and append
            nvcc_flags = extra_compile_args["nvcc"]
219
            extra_compile_args["nvcc"] = [f for f in nvcc_flags if not ("-O" in f or "-g" in f)]
220
221
222
            extra_compile_args["nvcc"].append("-O0")
            extra_compile_args["nvcc"].append("-g")

223
224
    sources = [os.path.join(extensions_dir, s) for s in sources]

225
    include_dirs = [extensions_dir]
226
227
228

    ext_modules = [
        extension(
229
            "torchvision._C",
230
            sorted(sources),
231
232
            include_dirs=include_dirs,
            define_macros=define_macros,
Soumith Chintala's avatar
Soumith Chintala committed
233
            extra_compile_args=extra_compile_args,
234
        )
235
    ]
236

237
    # ------------------- Torchvision extra extensions ------------------------
238
239
240
241
    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 []
242
243
244
245
246
247
248
249
250
    include_dirs += vision_include
    library_dirs = vision_library

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

251
252
253
    if sys.platform == "win32":
        image_macros += [("USE_PYTHON", None)]

254
    # Locating libPNG
255
256
    libpng = shutil.which("libpng-config")
    pngfix = shutil.which("pngfix")
257
    png_found = libpng is not None or pngfix is not None
258
259
260
261

    use_png = use_png and png_found
    if use_png:
        print("Found PNG library")
262
263
        if libpng is not None:
            # Linux / Mac
264
            min_version = "1.6.0"
265
266
            png_version = subprocess.run([libpng, "--version"], stdout=subprocess.PIPE)
            png_version = png_version.stdout.strip().decode("utf-8")
267
            png_version = parse_version(png_version)
268
            if png_version >= parse_version(min_version):
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
                image_include += [png_include]
278
                image_link_flags.append("png")
279
280
                print(f"  libpng version: {png_version}")
                print(f"  libpng include path: {png_include}")
281
            else:
282
283
284
                print("Could not add PNG image support to torchvision:")
                print(f"  libpng minimum version {min_version}, found {png_version}")
                use_png = False
285
286
        else:
            # Windows
287
288
            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")
289
290
            image_library += [png_lib]
            image_include += [png_include]
291
            image_link_flags.append("libpng")
292
293
294
    else:
        print("Building torchvision without PNG image support")
    image_macros += [("PNG_FOUND", str(int(use_png)))]
295

296
    # Locating libjpeg
297
    (jpeg_found, jpeg_conda, jpeg_include, jpeg_lib) = find_library("jpeglib", vision_include)
298

299
300
    use_jpeg = use_jpeg and jpeg_found
    if use_jpeg:
301
302
        print("Building torchvision with JPEG image support")
        image_link_flags.append("jpeg")
303
304
305
        if jpeg_conda:
            image_library += [jpeg_lib]
            image_include += [jpeg_include]
306
307
308
    else:
        print("Building torchvision without JPEG image support")
    image_macros += [("JPEG_FOUND", str(int(use_jpeg)))]
309

310
311
312
    # Locating nvjpeg
    # Should be included in CUDA_HOME for CUDA >= 10.1, which is the minimum version we have in the CI
    nvjpeg_found = (
313
314
315
        extension is CUDAExtension
        and CUDA_HOME is not None
        and os.path.exists(os.path.join(CUDA_HOME, "include", "nvjpeg.h"))
316
317
    )

318
319
    use_nvjpeg = use_nvjpeg and nvjpeg_found
    if use_nvjpeg:
320
321
        print("Building torchvision with NVJPEG image support")
        image_link_flags.append("nvjpeg")
322
323
324
    else:
        print("Building torchvision without NVJPEG image support")
    image_macros += [("NVJPEG_FOUND", str(int(use_nvjpeg)))]
325
326
327
328
329
330
331

    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"))
    )
332

333
    if use_png or use_jpeg:
334
335
336
337
338
339
340
341
342
343
344
345
        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,
            )
        )

346
    # Locating ffmpeg
347
    ffmpeg_exe = shutil.which("ffmpeg")
348
    has_ffmpeg = ffmpeg_exe is not None
349
    ffmpeg_version = None
350
351
352
353
    # 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
354
    if sys.platform != "linux" or (sys.version_info.major == 3 and sys.version_info.minor == 9):
355
        has_ffmpeg = False
356
357
    if has_ffmpeg:
        try:
358
            # This is to check if ffmpeg is installed properly.
359
            ffmpeg_version = subprocess.check_output(["ffmpeg", "-version"])
360
        except subprocess.CalledProcessError:
361
362
            print("Building torchvision without ffmpeg support")
            print("  Error fetching ffmpeg version, ignoring ffmpeg.")
363
364
            has_ffmpeg = False

365
    use_ffmpeg = use_ffmpeg and has_ffmpeg
366

367
    if use_ffmpeg:
368
        ffmpeg_libraries = {"libavcodec", "libavformat", "libavutil", "libswresample", "libswscale"}
369

370
371
        ffmpeg_bin = os.path.dirname(ffmpeg_exe)
        ffmpeg_root = os.path.dirname(ffmpeg_bin)
372
373
        ffmpeg_include_dir = os.path.join(ffmpeg_root, "include")
        ffmpeg_library_dir = os.path.join(ffmpeg_root, "lib")
374

375
        gcc = os.environ.get("CC", shutil.which("gcc"))
376
377
        platform_tag = subprocess.run([gcc, "-print-multiarch"], stdout=subprocess.PIPE)
        platform_tag = platform_tag.stdout.strip().decode("utf-8")
378
379
380

        if platform_tag:
            # Most probably a Debian-based distribution
381
382
            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)]
383
384
385
386
387
388
389
        else:
            ffmpeg_include_dir = [ffmpeg_include_dir]
            ffmpeg_library_dir = [ffmpeg_library_dir]

        for library in ffmpeg_libraries:
            library_found = False
            for search_path in ffmpeg_include_dir + include_dirs:
390
                full_path = os.path.join(search_path, library, "*.h")
391
392
393
                library_found |= len(glob.glob(full_path)) > 0

            if not library_found:
394
395
396
397
398
                print("Building torchvision without ffmpeg support")
                print(f"  {library} header files were not found, disabling ffmpeg support")
                use_ffmpeg = False
    else:
        print("Building torchvision without ffmpeg support")
399

400
401
402
403
404
    if use_ffmpeg:
        print("Building torchvision with ffmpeg support")
        print(f"  ffmpeg version: {ffmpeg_version}")
        print(f"  ffmpeg include path: {ffmpeg_include_dir}")
        print(f"  ffmpeg library_dir: {ffmpeg_library_dir}")
405
406

        # TorchVision base decoder + video reader
407
        video_reader_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "video_reader")
408
        video_reader_src = glob.glob(os.path.join(video_reader_src_dir, "*.cpp"))
409
410
        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"))
411
        # Torchvision video API
412
        videoapi_src_dir = os.path.join(this_dir, "torchvision", "csrc", "io", "video")
413
        videoapi_src = glob.glob(os.path.join(videoapi_src_dir, "*.cpp"))
414
        # exclude tests
415
        base_decoder_src = [x for x in base_decoder_src if "_test.cpp" not in x]
416

417
        combined_src = video_reader_src + base_decoder_src + videoapi_src
418

419
420
        ext_modules.append(
            CppExtension(
421
                "torchvision.video_reader",
422
                combined_src,
423
                include_dirs=[
424
                    base_decoder_src_dir,
425
                    video_reader_src_dir,
426
                    videoapi_src_dir,
427
                    extensions_dir,
428
                    *ffmpeg_include_dir,
429
                    *include_dirs,
430
                ],
431
                library_dirs=ffmpeg_library_dir + library_dirs,
432
                libraries=[
433
434
435
436
437
                    "avcodec",
                    "avformat",
                    "avutil",
                    "swresample",
                    "swscale",
438
                ],
439
440
                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"],
441
442
            )
        )
443

Prabhat Roy's avatar
Prabhat Roy committed
444
445
446
447
448
449
450
451
452
453
454
455
    # Locating video codec
    # CUDA_HOME should be set to the cuda root directory.
    # TORCHVISION_INCLUDE and TORCHVISION_LIBRARY should include the location to
    # video codec header files and libraries respectively.
    video_codec_found = (
        extension is CUDAExtension
        and CUDA_HOME is not None
        and any([os.path.exists(os.path.join(folder, "cuviddec.h")) for folder in vision_include])
        and any([os.path.exists(os.path.join(folder, "nvcuvid.h")) for folder in vision_include])
        and any([os.path.exists(os.path.join(folder, "libnvcuvid.so")) for folder in library_dirs])
    )

456
    use_video_codec = use_video_codec and video_codec_found
Prabhat Roy's avatar
Prabhat Roy committed
457
    if (
458
459
        use_video_codec
        and use_ffmpeg
Prabhat Roy's avatar
Prabhat Roy committed
460
461
        and any([os.path.exists(os.path.join(folder, "libavcodec", "bsf.h")) for folder in ffmpeg_include_dir])
    ):
462
        print("Building torchvision with video codec support")
Prabhat Roy's avatar
Prabhat Roy committed
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
        gpu_decoder_path = os.path.join(extensions_dir, "io", "decoder", "gpu")
        gpu_decoder_src = glob.glob(os.path.join(gpu_decoder_path, "*.cpp"))
        cuda_libs = os.path.join(CUDA_HOME, "lib64")
        cuda_inc = os.path.join(CUDA_HOME, "include")

        ext_modules.append(
            extension(
                "torchvision.Decoder",
                gpu_decoder_src,
                include_dirs=include_dirs + [gpu_decoder_path] + [cuda_inc] + ffmpeg_include_dir,
                library_dirs=ffmpeg_library_dir + library_dirs + [cuda_libs],
                libraries=[
                    "avcodec",
                    "avformat",
                    "avutil",
                    "swresample",
                    "swscale",
                    "nvcuvid",
                    "cuda",
                    "cudart",
                    "z",
                    "pthread",
                    "dl",
486
                    "nppicc",
Prabhat Roy's avatar
Prabhat Roy committed
487
488
489
490
491
                ],
                extra_compile_args=extra_compile_args,
            )
        )
    else:
492
493
494
495
496
497
498
499
500
501
502
        print("Building torchvision without video codec support")
        if (
            use_video_codec
            and use_ffmpeg
            and not any([os.path.exists(os.path.join(folder, "libavcodec", "bsf.h")) for folder in ffmpeg_include_dir])
        ):
            print(
                "  The installed version of ffmpeg is missing the header file 'bsf.h' which is "
                "  required for GPU video decoding. Please install the latest ffmpeg from conda-forge channel:"
                "   `conda install -c conda-forge ffmpeg`."
            )
Prabhat Roy's avatar
Prabhat Roy committed
503

504
505
506
507
508
    return ext_modules


class clean(distutils.command.clean.clean):
    def run(self):
509
        with open(".gitignore") as f:
510
            ignores = f.read()
511
            for wildcard in filter(None, ignores.split("\n")):
512
513
514
515
516
517
518
519
520
521
                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)


522
if __name__ == "__main__":
523
    print(f"Building wheel {package_name}-{version}")
524
525
526

    write_version_file()

527
    with open("README.rst") as f:
528
529
530
531
532
533
        readme = f.read()

    setup(
        # Metadata
        name=package_name,
        version=version,
534
535
536
537
        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",
538
        long_description=readme,
539
        license="BSD",
540
        # Package info
541
        packages=find_packages(exclude=("test",)),
542
        package_data={package_name: ["*.dll", "*.dylib", "*.so"]},
543
544
545
546
547
548
        zip_safe=False,
        install_requires=requirements,
        extras_require={
            "scipy": ["scipy"],
        },
        ext_modules=get_extensions(),
549
        python_requires=">=3.7",
550
        cmdclass={
551
552
553
            "build_ext": BuildExtension.with_options(no_python_abi_suffix=True),
            "clean": clean,
        },
554
    )