Python 调用 C++ 模块通常有三种方法:
1. 直接调用动态库:通过系统原生的 .so(Linux)或 .dll(Windows)文件进行调用。
2. 使用 pybind11:借助这一强大的 C++/Python 绑定库,可以便捷地导出 C++ 函数和类。
3. 使用 PyTorch Extension:适用于在 PyTorch 框架中集成自定义 C++ 算子。
下面分别展示在 Linux 和 Windows 系统下,使用第一种方法、第二种方法以及 PyTorch Extension 的具体实现示例。
1. Linux下直接调用so:
example.cpp:
// example.cpp extern "C" {int add(int a, int b) {return a + b;}double multiply(double a, double b) {return a * b;} }
编译指令:
g++ -shared -fPIC -o libexample.so example.cpp
main.py:
from ctypes import cdll, c_int, c_double# 加载DLL lib = cdll.LoadLibrary('./libexample.so')# 调用函数lib.add.argtypes = [c_int, c_int] lib.add.restype = c_int result = lib.add(3, 4) print(f"3 + 4 = {result}")# 对于有明确参数类型的函数,可以指定类型 lib.multiply.argtypes = [c_double, c_double] lib.multiply.restype = c_double result = lib.multiply(2.5, 4.0) print(f"2.5 * 4.0 = {result}")
2. Windows下直接调用dll:
example.cpp:
// 使用__declspec(dllexport)导出函数 extern "C" {__declspec(dllexport) int add(int a, int b) {return a + b;}__declspec(dllexport) double multiply(double a, double b) {return a * b;} }
编译指令:
cl /LD example.cpp /link /OUT:example.dll
main.py:
import ctypes from ctypes import cdll, c_int, c_double# 加载DLL lib = cdll.LoadLibrary('./example.dll')# 调用函数 lib.add.argtypes = [c_int, c_int] lib.add.restype = c_int result = lib.add(3, 4) print(f"3 + 4 = {result}")# 对于有明确参数类型的函数,可以指定类型 lib.multiply.argtypes = [c_double, c_double] lib.multiply.restype = c_double result = lib.multiply(2.5, 4.0) print(f"2.5 * 4.0 = {result}")
3. Linux下使用pybind调用:
example_bind.cpp:
#include <pybind11/pybind11.h>int add(int a, int b) {return a + b; }double multiply(double a, double b) {return a * b; }double divv(double a, double b) {return (a / b); }// 创建Python模块 PYBIND11_MODULE(example_bind, m) {m.def("add", &add, "A function that adds two numbers");m.def("multiply", &multiply, "A function that multiplies two numbers");m.def("divv", &divv, "A function that div two numbers");}
编译指令:
g++ -O3 -Wall -shared -std=c++11 -fPIC $(python3 -m pybind11 --includes) example_bind.cpp -o example_bind$(python3-config --extension-suffix) 或: python setup.py build_ext --inplace
setup.py:
# setup.py from setuptools import setup, Extension import pybind11ext_modules = [Extension('example_bind',['example_bind.cpp'],include_dirs=[pybind11.get_include()],language='c++'), ]setup(name='example_bind',ext_modules=ext_modules,zip_safe=False, )
main.py:
import example_bindresult = example_bind.add(3, 4) print(f"3 + 4 = {result}")result = example_bind.multiply(2.5, 4.0) print(f"2.5 * 4.0 = {result}")result = example_bind.divv(2.5, 4.0) print(f"2.5 / 4.0 = {result}")
4. Windows下使用pybind调用:
example_bind.cpp:
#include <pybind11/pybind11.h>int add(int a, int b) {return a + b; }double multiply(double a, double b) {return a * b; }double divv(double a, double b) {return (a / b); }// 创建Python模块 PYBIND11_MODULE(example_bind, m) {m.def("add", &add, "A function that adds two numbers");m.def("multiply", &multiply, "A function that multiplies two numbers");m.def("divv", &divv, "A function that div two numbers");}
编译指令:
cl /EHsc /LD /I C:/Users/xx/AppData/Local/Programs/Python/Python39/include /I C:/Users/xx/AppData/Local/Programs/Python/Python39/Lib/site-packages/pybind11/include example_bind.cpp /link /LIBPATH:C:/Users/xx/AppData/Local/Programs/Python/Python39/libs python39.lib /OUT:example_bind.pyd
或:
python setup.py build_ext --inplace
setup.py:
# setup.py from setuptools import setup, Extension import pybind11ext_modules = [Extension('example_bind',['example_bind.cpp'],include_dirs=[pybind11.get_include()],language='c++'), ]setup(name='example_bind',ext_modules=ext_modules,zip_safe=False, )
main.py:
import example_bindresult = example_bind.add(3, 4) print(f"3 + 4 = {result}")result = example_bind.multiply(2.5, 4.0) print(f"2.5 * 4.0 = {result}")result = example_bind.divv(2.5, 4.0) print(f"2.5 / 4.0 = {result}")
5. 使用 PyTorch Extension 调用:
torch_extension.cpp:
#include <torch/extension.h>torch::Tensor add_tensors(torch::Tensor a, torch::Tensor b) {return a + b; }// 注册扩展 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("add_tensors", &add_tensors, "Add two tensors"); }
setup.py:
# setup.py from setuptools import setup from torch.utils.cpp_extension import CppExtension, BuildExtensionsetup(name='torch_extension',ext_modules=[CppExtension('torch_extension', ['torch_extension.cpp'])],cmdclass={'build_ext': BuildExtension} )
编译指令:
python setup.py install
main.py:
import torch import torch_extensiona = torch.rand(100)print(a) b = torch.rand(100) result = torch_extension.add_tensors(a, b) print(result)