526互联

Pybind11：使用C++编写Python模块

发布时间 2023-08-23 11:18:38作者: undermyth

放假摆了一周了。看论文实在不是什么有意思的活。

这两天研究了一下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带来的调用开销可能是比较大的，导致反复调用的开销还是比较大。

当然，这只是个人解释。后面说不定还有什么妙の原因，之后再探索吧。