pybind11&C++环境下如何实现vector引用传递并同步数据到Python侧

问题根因 #

当前问题是pybind11默认的std::vector类型转换规则导致的：当Python侧的list传入绑定的C接口时，pybind11会拷贝生成一个独立的C std::vector对象，C++侧修改的是这个拷贝后的对象，不会同步到Python侧原来的list上。

方案1：修改接口为返回vector（代码改动最小） #

直接把C++的calc方法调整为返回填充好的std::vector<double>，pybind11会自动把返回的vector转为Python的list，不需要修改绑定逻辑，改动量最小。

C++代码调整 #

class Calculator
{
public:
    Calculator() {}
    // 改返回值为std::vector<double>，去掉入参里的results
    static std::vector<double> calc(const Setup& setup) {
        std::vector<double> results;
        // 原有填充results的逻辑保持不变
        return results;
    }
};

// 绑定部分无需修改，保持原有逻辑即可

Python侧调用调整 #

results = a_calculator.calc(a_setup)

该方案优点是无需处理类型转换细节，代码量最少，适合对性能要求不高的场景。

方案2：用numpy数组实现零拷贝引用传递（性能最优） #

如果数据量较大，不想产生拷贝开销，可以直接使用pybind11的numpy数组接口，直接操作Python侧numpy数组的内存，改动很小。

C++代码调整 #

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

class Calculator
{
public:
    Calculator() {}
    // 把results参数改为py::array_t<double>类型接收numpy数组
    static void calc(const Setup& setup, py::array_t<double> results) {
        // 获取numpy数组的可写指针
        auto buf = results.mutable_unchecked<1>();
        // 直接操作内存填充数据，示例如下
        for (py::ssize_t i = 0; i < buf.shape(0); i++) {
            buf(i) = 待填充的数值;
        }
    }
};

// 绑定逻辑保持原有内容不变即可

Python侧调用调整 #

# 提前初始化对应长度、类型为float64的numpy数组
results = np.zeros(所需数据长度, dtype=np.float64)
a_calculator.calc(a_setup, results)
# 此时results已被C++侧修改完成，可直接使用

该方案为零拷贝实现，性能远高于方案1，适合处理大数据量场景。

可选方案：绑定std::vector实现引用传递 #

如果必须保留C++侧std::vector<double>&的参数类型，可以在绑定代码中增加std::vector的绑定声明：

PYBIND11_MODULE(one_calculator, m) {
    // 增加该行绑定std::vector<double>为Python侧的VectorDouble类型
    py::bind_vector<std::vector<double>>(m, "VectorDouble");

    // 原有绑定逻辑保持不变
    m.doc() = "pybind11 one_calculator plugin";
    py::class_<Setup>(m, "Setup")
        .def(py::init<int>());
    py::class_<Calculator>(m, "Calculator")
        .def(py::init<>())
        .def("calc", &Calculator::calc);
}

Python侧调用时使用绑定的VectorDouble类型即可：

from one_calculator import Setup, Calculator, VectorDouble

results = VectorDouble()
a_calculator.calc(a_setup, results)
# 可根据需要转为原生list或numpy数组
list_results = list(results)
np_results = np.array(results)

内容的提问来源于stack exchange，提问作者xycs