pytorch3d/setup.py
Christoph Lassner 039e02601d examples and docs.
Summary: This diff updates the documentation and tutorials with information about the new pulsar backend. For more information about the pulsar backend, see the release notes and the paper (https://arxiv.org/abs/2004.07484). For information on how to use the backend, see the point cloud rendering notebook and the examples in the folder docs/examples.

Reviewed By: nikhilaravi

Differential Revision: D24498129

fbshipit-source-id: e312b0169a72b13590df6e4db36bfe6190d742f9
2020-11-03 13:06:35 -08:00

112 lines
3.7 KiB
Python
Executable File

#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import glob
import os
import runpy
import warnings
import torch
from setuptools import find_packages, setup
from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
def get_extensions():
this_dir = os.path.dirname(os.path.abspath(__file__))
extensions_dir = os.path.join(this_dir, "pytorch3d", "csrc")
sources = glob.glob(os.path.join(extensions_dir, "**", "*.cpp"), recursive=True)
source_cuda = glob.glob(os.path.join(extensions_dir, "**", "*.cu"), recursive=True)
extension = CppExtension
extra_compile_args = {"cxx": ["-std=c++14"]}
define_macros = []
force_cuda = os.getenv("FORCE_CUDA", "0") == "1"
if (torch.cuda.is_available() and CUDA_HOME is not None) or force_cuda:
extension = CUDAExtension
sources += source_cuda
define_macros += [("WITH_CUDA", None)]
cub_home = os.environ.get("CUB_HOME", None)
nvcc_args = [
"-std=c++14",
"-DCUDA_HAS_FP16=1",
"-D__CUDA_NO_HALF_OPERATORS__",
"-D__CUDA_NO_HALF_CONVERSIONS__",
"-D__CUDA_NO_HALF2_OPERATORS__",
]
if cub_home is None:
warnings.warn(
"The environment variable `CUB_HOME` was not found. "
"NVIDIA CUB is required for compilation and can be downloaded "
"from `https://github.com/NVIDIA/cub/releases`. You can unpack "
"it to a location of your choice and set the environment variable "
"`CUB_HOME` to the folder containing the `CMakeListst.txt` file."
)
else:
nvcc_args.insert(
0, "-I%s" % (os.path.realpath(cub_home).replace("\\ ", " "))
)
nvcc_flags_env = os.getenv("NVCC_FLAGS", "")
if nvcc_flags_env != "":
nvcc_args.extend(nvcc_flags_env.split(" "))
# It's better if pytorch can do this by default ..
CC = os.environ.get("CC", None)
if CC is not None:
CC_arg = "-ccbin={}".format(CC)
if CC_arg not in nvcc_args:
if any(arg.startswith("-ccbin") for arg in nvcc_args):
raise ValueError("Inconsistent ccbins")
nvcc_args.append(CC_arg)
extra_compile_args["nvcc"] = nvcc_args
sources = [os.path.join(extensions_dir, s) for s in sources]
include_dirs = [extensions_dir]
ext_modules = [
extension(
"pytorch3d._C",
sources,
include_dirs=include_dirs,
define_macros=define_macros,
extra_compile_args=extra_compile_args,
)
]
return ext_modules
# Retrieve __version__ from the package.
__version__ = runpy.run_path("pytorch3d/__init__.py")["__version__"]
if os.getenv("PYTORCH3D_NO_NINJA", "0") == "1":
class BuildExtension(torch.utils.cpp_extension.BuildExtension):
def __init__(self, *args, **kwargs):
super().__init__(use_ninja=False, *args, **kwargs)
else:
BuildExtension = torch.utils.cpp_extension.BuildExtension
setup(
name="pytorch3d",
version=__version__,
author="FAIR",
url="https://github.com/facebookresearch/pytorch3d",
description="PyTorch3D is FAIR's library of reusable components "
"for deep Learning with 3D data.",
packages=find_packages(exclude=("configs", "tests")),
install_requires=["torchvision>=0.4", "fvcore"],
extras_require={
"all": ["matplotlib", "tqdm>4.29.0", "imageio", "ipywidgets"],
"dev": ["flake8", "isort", "black==19.3b0"],
},
ext_modules=get_extensions(),
cmdclass={"build_ext": BuildExtension},
)