pytorch3d/setup.py
Christoph Lassner 960fd6d8b6 pulsar interface unification.
Summary:
This diff builds on top of the `pulsar integration` diff to provide a unified interface for the existing PyTorch3D point renderer and Pulsar. 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.

The unified interfaces are completely consistent. Switching the render backend is as easy as using `renderer = PulsarPointsRenderer(rasterizer=rasterizer).to(device)` instead of `renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor)` and adding the `gamma` parameter to the forward function. All PyTorch3D camera types are supported as far as possible; keyword arguments are properly forwarded to the camera. The `PerspectiveCamera` and `OrthographicCamera` require znear and zfar as additional parameters for the forward pass.

Reviewed By: nikhilaravi

Differential Revision: D21421443

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

108 lines
3.6 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 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)
if cub_home is None:
raise Exception(
"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."
)
nvcc_args = [
"-I%s" % (os.path.realpath(cub_home).replace("\\ ", " ")),
"-std=c++14",
"-DCUDA_HAS_FP16=1",
"-D__CUDA_NO_HALF_OPERATORS__",
"-D__CUDA_NO_HALF_CONVERSIONS__",
"-D__CUDA_NO_HALF2_OPERATORS__",
]
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},
)