c++ 如何将torch路径添加到Pybind11扩展?

xyhw6mcr  于 2023-06-25  发布在  其他
关注(0)|答案(1)|浏览(127)

我正在用C++编写一个python扩展,使用pybind11,我正在尝试创建setup.py文件,这是我目前为止所做的:

  1. from glob import glob
  2. import setuptools
  3. from pybind11.setup_helpers import Pybind11Extension, build_ext
  4. ext_modules = [
  5. Pybind11Extension(
  6. "my_ext",
  7. sorted(glob("src/my_ext/**/*.cc")),
  8. )
  9. ]
  10. setuptools.setup(
  11. name="my_ext",
  12. version="0.1",
  13. package_dir={"": "src/my_ext/python/"},
  14. cmdclass={"build_ext": build_ext},
  15. ext_modules=ext_modules,
  16. )

但是,当我运行pip install .时,我得到了这个错误:

  1. In file included from src/my_ext/cc/thing.cc:7:
  2. src/my_ext/cc/thing.h:9:10: fatal error: 'torch/torch.h' file not found
  3. #include <torch/torch.h>
  4. ^~~~~~~~~~~~~~~
  5. 1 error generated.
  6. error: command '/usr/bin/clang' failed with exit code 1

有没有什么参数可以传递给Pybind11Extension,让它找到torch并成功构建?

dwthyt8l

dwthyt8l1#

我自己想出来的。只需要使用一些pytorch的helper函数,并查看torch.utils.cpp_extension.CppExtension以获得指导:

  1. import os
  2. from glob import glob
  3. import setuptools
  4. from pybind11.setup_helpers import Pybind11Extension, build_ext
  5. from torch.utils.cpp_extension import include_paths, library_paths
  6. include_dirs = include_paths()
  7. torch_library_paths = library_paths()
  8. libraries = ["c10", "torch", "torch_cpu", "torch_python"]
  9. ext_modules = [
  10. Pybind11Extension(
  11. "my_ext",
  12. sorted(glob("src/my_ext/**/*.cc")),
  13. include_dirs=include_dirs,
  14. libraries=libraries,
  15. library_dirs=torch_library_paths,
  16. )
  17. ]
  18. setuptools.setup(
  19. name="my_ext",
  20. version="0.1",
  21. cmdclass={"build_ext": build_ext},
  22. ext_modules=ext_modules,
  23. )
展开查看全部

相关问题