pybind11 Eigen + STL Containers + Numpy

弓宏茂

2023-12-01

Pybind11

梳理一下 Pybind11 中 Eigen 和 STL 容器与 Python 交互的方式。需要先了解Functions章节中关于返回值和调用的规则。

Smart Pointers

STL 容器

Pybind11 已经自动支持 std::vector<>/std::deque<>/std::list<>/std::array<>/std::valarray<>, std::set<>/std::unordered_set<>, and std::map<>/std::unordered_map<> 和 Python list, set and dict 之间的转换。引用头文件 pybind11/pybind11.h 即可。如果想使用全部类，可以引用 pybind11/stl.h

Arbitrary nesting of any of these types is possible. 比如 std::vector<std::vector<int>>

问题

该方法的缺点在于每一次在 C++ 和 Python 之间转换时都要隐式转换，并且制作一份拷贝。所以下面的例子会失效

# Example 1
# C++
void append_1(std::vector<int> &v) {
   v.push_back(1);
}

# python
>>> v = [5, 6]
>>> append_1(v)
>>> print(v)
[5, 6]

# Example 2
# C++
/* ... definition ... */

class MyClass {
    std::vector<int> contents;
};

/* ... binding code ... */

py::class_<MyClass>(m, "MyClass")
    .def(py::init<>())
    .def_readwrite("contents", &MyClass::contents);
    
# Python
>>> m = MyClass()
>>> m.contents = [5, 6]
>>> print(m.contents)
[5, 6]
>>> m.contents.append(7)
>>> print(m.contents)
[5, 6]

上述例子 1 中，Python 向 C++ 传递 List 后，C++ 端制作了一份拷贝，所以两者使用的是不同的数据；例子 2 中，虽然pybind11将类的属性绑定到 Python，在 Python 端能够直接访问，但是例如 append 方法无法正常使用也是相同的原因。（Python 侧和 Cpp 侧并不是相同的实例）

Opaque 类

为了解决上述问题，Pybind11 提供了 PYBIND11_MAKE_OPAQUE(T) 宏定义将其变为 Opaque Type。Opaque Types 需要有对应的 class_ 声明。例如：

// Opaque Type declaration
PYBIND11_MAKE_OPAQUE(std::vector<int>);

// Pybind11 Class
// 定义其需要暴露给 Python 的方法
py::class_<std::vector<int>>(m, "IntVector")
    .def(py::init<>())
    .def("clear", &std::vector<int>::clear)
    .def("pop_back", &std::vector<int>::pop_back)
    .def("__len__", [](const std::vector<int> &v) { return v.size(); })
    .def("__iter__", [](std::vector<int> &v) {
       return py::make_iterator(v.begin(), v.end());
    }, py::keep_alive<0, 1>()) /* Keep vector alive while iterator is used */
    // ....

补充–函数指针

补充 C++ 中函数指针，参考，声明函数指针

data_types (*func_pointer)(data_types arg1, data_types arg2, ...,data_types argn);

比如

int test(int a);

// 声明函数指针 fp
int (* fp) (int); 
// 赋值
fp = test;

指向类成员函数的函数指针。需要注意两点

函数指针赋值时需要使用 &
调用函数时，需要使用 .* 或者 ->*

例如

void (A::* ptr)(int) = &A::setA; // member function pointer

// call
pa->*ptr(100);  // pa 是 A 的指针
a.*ptr(100);    //  a 是 A 的实例

更多示例

我们将之前的 std::vector<int> 继续扩展，添加 append 和索引支持。

// Opaque Type declaration
PYBIND11_MAKE_OPAQUE(std::vector<int>);

// Pybind11 Class
// 定义其需要暴露给 Python 的方法
py::class_<std::vector<int>>(m, "IntVector")
    .def(py::init<>())
    .def("clear", &std::vector<int>::clear)
    .def("pop_back", &std::vector<int>::pop_back)
    .def("__len__", [](const std::vector<int> &v) { return v.size(); })
    .def("__iter__", [](std::vector<int> &v) {
       return py::make_iterator(v.begin(), v.end());
    }, py::keep_alive<0, 1>()) /* Keep vector alive while iterator is used */
    // append function for list 
    .def("append", (void (std::vector<int>::*)(const int&)) &std::vector<int>::push_back, 
                   py::keep_alive<1, 2>())
    // support index
    .def("__getitem__", (int& (std::vector<int>::*)(size_t)) &std::vector<int>::operator[], py::return_value_policy::reference_internal)

append 方法中使用 keep_alive 规则。当 C++ 侧的对象是任何一种容器的时候需要使用该规则。keep_alive<Nurse, Patient> 表明 Patient 所指明的参数至少应该和 Nurse 所指的参数的生命周期一样长。（当加入元素到容器中时，如果不适用 keep alive，当 Python 侧的变量释放时会自动释放此元素，但是 C++ 仍然将其存放在容器中，会导致段错误）0 代表返回值，1 一般代表 this，其余按传入参数依次排列。

由于 Lambda 函数里面不让用 this 指针，自定义构造函数（接受 Python 类）用 Lambda 函数有些困难。另一种方法可以是自己写一个类继承 std::vector 并自己写构造函数然后通过 Pybind 绑定到 Python 上。

Eigen

Pybind11 支持 Eigen 类型转换。引用 pybind11/eigen.h 即可。

Python 到 C++

当 Python 向 C++ 传入 np.ndarray 时，Pybind11 会将其拷贝到 Eigen Dense Matrix 临时变量中，并将临时变量传入调用的函数中。

如果需要传入引用值，则只需要使用 Eigen::Ref<const MatrixType> 或者 Eigen::Ref<MatrixType> 即可，这需要传入的 ndarray 与 MatrixType 类型匹配并且存储格式（列主 / 行主）一致。但是对于后者，还需要传入数组是可写的 (a.flags['writeable'] = True)

C++ 到 Python

当 C++ 返回 Dense Matrix 时，返回的是原矩阵的引用。Numpy 并不拥有数据，且返回对象的生命周期绑定到返回的 array 上。如果返回值前加上 const 关键字，则在 Python 端不能修改数据。

当返回引用或者指针时，适用 Pybind11 automatic 规则。如果是指针，则 Python 负责调用所指对象的析构函数，C++ 侧不需要。如果是引用，则做一份拷贝。

如果返回的是 Map / Ref / Block ，Pybind11 默认传回的是引用，因此需要保证 Python 访问数据时，数据仍然有效。安全的方法是将使用 copy 规则。或者使用 reference_internal / keep_alive 保证返回数据和传入数据或者其对象用相同的生命周期。

Numpy

如果单纯对 Numpy 数组进行数据操作，最方便的还是使用上面 Eigen 方法，能够直接用 C++ 处理。Pybind11 同时还提供了 py::array_t 用于接受 Numpy Array 参数。

void f(py::array_t<double, py::array::c_style | py::array::forcecast> array);

以上这个例子是接收 C 存储格式的 Numpy Array 。py::array::forcecast 是保证传入参数会被尝试转换为当前的 Array 类型而不是使用其他重载函数。

传入 array 使用：

.ndim() 返回张量的维数；.shape() 返回 ssize_t * 指向存储各维度大小的数组。
读取数据。分为两组：.at(i,j,k) .data(i, j, k) 与 .mutable_at(i,j,k) .mutable_data(i, j, k) 。前者返回常（const）引用与指针，后者则是可以修改数据。
所有的数据内存中相邻。
省略号

py::array a = /* A NumPy array */;

// 相当于 a[0, ..., 0]
py::array b = a[py::make_tuple(0, py::ellipsis(), 0)];

Numpy Dtype

能够拓展到更多的数据类型。

struct A {
    int x;
    double y;
};

struct B {
    int z;
    A a;
};

// ...
PYBIND11_MODULE(test, m) {
    // ...

    PYBIND11_NUMPY_DTYPE(A, x, y);
    PYBIND11_NUMPY_DTYPE(B, z, a);
    /* now both A and B can be used as template arguments to py::array_t */
}