PointMLP/pointnet2_ops_lib/setup.py

40 lines
1.2 KiB
Python
Raw Normal View History

2021-10-04 07:25:18 +00:00
import glob
import os
import os.path as osp
from setuptools import find_packages, setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
this_dir = osp.dirname(osp.abspath(__file__))
_ext_src_root = osp.join("pointnet2_ops", "_ext-src")
_ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob(
osp.join(_ext_src_root, "src", "*.cu")
)
_ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*"))
requirements = ["torch>=1.4"]
exec(open(osp.join("pointnet2_ops", "_version.py")).read())
os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;5.0;6.0;6.1;6.2;7.0;7.5"
setup(
name="pointnet2_ops",
version=__version__,
author="Erik Wijmans",
packages=find_packages(),
install_requires=requirements,
ext_modules=[
CUDAExtension(
name="pointnet2_ops._ext",
sources=_ext_sources,
extra_compile_args={
"cxx": ["-O3"],
"nvcc": ["-O3", "-Xfatbin", "-compress-all"],
},
include_dirs=[osp.join(this_dir, _ext_src_root, "include")],
)
],
cmdclass={"build_ext": BuildExtension},
include_package_data=True,
)