使用pytorch的pybind11实现c++扩展的简单调用
1、首先使用c++实现需要的功能,以简单加法和乘法为例:
// 本文件名为test.cpp
#include <iostream>
#include <torch/torch.h>
using namespace std;
int add(int i, int j){
return i + j;
}
int mul(int i, int j){
return i * j;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
m.def("add", &add, "int add"); // 第一个参数为python调用时的函数名,第二个参数为c++函数名,第三个为相应的描述
m.def("mul", &mul, "int multiply");
}
2、新建一个setup.py文件,实现:
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
setup(
name="test", # python中导入的名称
ext_modules=[
CUDAExtension("test", ["./test.cpp"]) # 若没有使用到cuda,也可以使用CppExtension, 此处的test名和上面name的需要一样
],
cmdclass={
'build_ext': BuildExtension
}
)
3、编译
在setup.py文件的同级目录下,终端输入:
python setup.py build_ext --inplace
其中–inplace是将生成的.so文件复制到当前目录下,也可以省略,但调用时就需要到生成的buile文件中导入
4、python调用c++扩展
import test
a = test.add(2, 3)
print(a) # 5
b = test.mul(2, 3)
print(b) # 6