Pybind11:使用C++编写Python模块
放假摆了一周了。看论文实在不是什么有意思的活。
这两天研究了一下Pybind11的用法。使用C/C++和Python混合编程的想法很早就有了,在大一的一次比赛时曾经实践过(虽然不是我写的),当时获得了比较显著的性能提升。但是当时用的是Swig,据队友说Swig对于NumPy的支持极为阴间,当时调试花了好几天的时间。在混合编程中NumPy的传递极为重要,因为混合编程的主要使用场景就在于python不擅长进行大规模的数值计算,以及不擅长进行并行,而python科学计算的核心就在于NumPy。
相比之下,Pybind11比Swig更新,而且对于NumPy有专门的支持。它是一组C++的头文件,功能类似于Boost.Python,但更加轻量化。
基本使用
代码
找到两篇很详细的Blog,一篇是知乎上的,一篇是一个个人博客。
基本原理是使用C++编译器将cpp模块生成动态库(.so/.pyd),python能够直接识别动态库为模块导入进行使用。
C++模块正常编写,包括函数或者类。在最后需要加上PYBIND11_MODULE
进行绑定。能够绑定的对象包括:函数、类(包括重载、继承、操作符重载、虚函数等)。C++操作符重载一部分直接对应于Python的操作符,还有一部分对应于Python的魔术方法。转换成魔术方法时,需要由C++侧去匹配魔术方法的参数表。参数表中的self
对应于一个该类型的引用。
其他的详细步骤我不写了,上面两篇博客写得很详细了。放一个我自己写的例子。这里有一部分代码是用CodeGeeX写的,不得不承认我写代码没有它快qwq
#include "pybind11/detail/common.h"
#include <pybind11/pybind11.h>
#include <pybind11/operators.h>
#include <omp.h>
#include <iostream>
#include <string>
template<class T>
class Vector {
public:
Vector(int size) {
this->size = size;
v = new T[size];
for (int i = 0; i < size; ++i) {
v[i] = 0;
}
}
~Vector() {
delete[] v;
}
T& operator[](int i) {
return v[i];
}
const Vector<T>& operator+=(const Vector<T>& other) {
if (size != other.size) {
throw "Vector sizes do not match";
}
#pragma omp parallel for
for (int i = 0; i < size; ++i) {
v[i] += other.v[i];
}
return *this;
}
void print() const{
for (int i = 0; i < size; ++i) {
std::cout << v[i] << " ";
}
std::cout << std::endl;
}
const std::string toString() const{
std::string s;
for (int i = 0; i < size; ++i) {
s += std::to_string(v[i]) + " ";
}
return s;
}
private:
T* v;
int size;
};
/* do binding */
PYBIND11_MODULE(vector, obj) {
pybind11::class_<Vector<double> > VecClass(obj, "Vector");
VecClass.def(pybind11::init<int>());
VecClass.def(pybind11::self += pybind11::self);
VecClass.def("__getitem__", &Vector<double>::operator[]);
VecClass.def("__setitem__", [](Vector<double>& v, int i, double x) {v[i] = x;});
VecClass.def("__str__", &Vector<double>::toString);
VecClass.def("__repr__", &Vector<double>::toString);
VecClass.def("print", &Vector<double>::print);
}
顺带一提,C++的模版在这里几乎起不到作用,因为Pybind11必须将模版实例化才能进行绑定,和没有模版几乎没区别。
编译
本人已经放弃CMake了,不能理解CMake的逻辑,还要记很多东西,不如直接Makefile。所以这里写的是直接按命令编译。
含Pybind11的C++代码编译命令是
$ c++ -O3 -Wall -shared -std=c++11 -fPIC $(python3 -m pybind11 --includes) example.cpp -o example$(python3-config --extension-suffix)
其实很简单对吧。标准需要在C++11以上,$(python3 -m pybind11 --includes)
是pybind11的头文件include路径。顺带一提,python-dev的include路径是$(python3-config --includes)
,链接库选项是$(python3-config --ldflags)
。$(python3-config --extension-suffix)
是与python版本和系统相关的后缀名,在我的电脑上以.so结尾。
特别提示,在Darwin(MacOS)上需要再加上-undefined dynamic_lookup
。这是一个很神奇的选项,我也不太理解它干什么用,好像是跳过编译阶段的undefined symbols,等到链接时再寻找。
运行
只要这个库在python脚本的同一目录或者PYTHONPATH
中,python就可以直接import它。模块名称是PYBIND11_MODULE
的第一个参数。
附加环节:使用pyi进行代码提示
使用C++生成的python库运行是没有问题的,但是不会有任何的代码补全和提示。主流IDE都是使用pyi文件来进行代码提示的,github上的项目pybind11-stubgen能够做到生成任意python模块的pyi文件。
在库文件的同一目录下运行pybind11-stubgen,比如对于上面的vector模块,执行
pybind11-stubgen vector --output-dir . --no-setup
它将会生成一个文件夹,里面包含了一个__init__.pyi
。这就是我们需要的东西。如果把--no-setup
去掉,还会生成一个setup.py
,但我不知道这有啥用。
要有条理地使用这个pyi文件,可以将库打包成一个包。也就是新建一个文件夹,把库文件和pyi塞进去,然后再加一个__init__.py
使其成为包。注意,__init__.py
中需要导入__init__.pyi
中__all__
里包括的变量才能使代码提示正常工作。
(可能也可以使用那个setup.py
)
踩坑
第一个坑:-undefined dynamic_lookup
。上面说过了。
第二个坑:最好使cpp的文件名、PYBIND11_MODULE
的模块名、库文件的前缀三者保持一致,cpp的类名/函数名和PYBIND11_MODULE
中绑定的python类名/函数名一致。其实不一致也没什么意义,但是不一致的话很容易在导入包的时候发生错误,最后还是不要这么做。
第三个坑:官方提供了操作符重载的简便写法
#include <pybind11/operators.h>
PYBIND11_MODULE(example, m) {
py::class_<Vector2>(m, "Vector2")
.def(py::init<float, float>())
.def(py::self + py::self)
.def(py::self += py::self)
.def(py::self *= float())
.def(float() * py::self)
.def(py::self * float())
.def(-py::self)
.def("__repr__", &Vector2::toString);
}
希望没有其他朋友像我一样漏看了第一行的头文件,导致py::self
找不到google了半个小时。
C++ vs NumPy
都写到这里了,我觉得可以解决一个我长时间以来不理解的问题:
Python比C++慢吗?慢多少?使用C++混合编程是否确实能提高速度?
我用上面的Vector类简单进行了一个测试。进行比较的操作符是运算后赋值(+=)。具体来说,是这样的:
from cpp_vec import Vector
import random
import time
import numpy as np
x = Vector(100000)
y = Vector(100000)
# assign each element in x with random double precision number
for i in range(100000):
x[i] = random.random()
y[i] = random.random()
# start timing
start = time.time()
for i in range(100):
x += y
end = time.time()
print("Time elapsed by cpp_vec: %f" % (end - start))
x = np.zeros(100000)
y = np.zeros(100000)
for i in range(100000):
x[i] = random.random()
y[i] = random.random()
start = time.time()
for i in range(100):
x += y
end = time.time()
print("Time elapsed by numpy: %f" % (end - start))
我在C++里开启了OpenMP(上面的代码里面写了),这是公平的,因为NumPy本身也是C实现的,它能够绕过GIL锁来使用多线程。上面的设置导致的结果是
Time elapsed by cpp_vec: 0.013195
Time elapsed by numpy: 0.004841
但如果将数组的大小开到10000000,循环次数取消(保证理论上的计算量一样),得到的结果是
Time elapsed by cpp_vec: 0.011262
Time elapsed by numpy: 0.016986
C++在大规模的代数运算上还是有优势的(何况这只是非常naive的一个实现)。但是Pybind11带来的调用开销可能是比较大的,导致反复调用的开销还是比较大。
当然,这只是个人解释。后面说不定还有什么妙の原因,之后再探索吧。