| 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, | |
| ) | |