使用pytorch实现python调用c++

使用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