无需复制即可将 C++ 的 Eigen::Matrix 数组返回到 Python

Question

我有一些生成和操作矩阵数组的 C++ 代码Eigen。最后我想在 python 中使用这些矩阵，并认为这可能是pybind11.基本上我想要在 python 中返回的是两个嵌套列表/numpy 数组 mat_a(I, 4, 4)和mat_b(J, K, 4, 4). 因为我必须在 C++ 中做很多线性代数的东西，所以我想使用 Eigen，我使用的数据结构是 std::array<std::array<Eigen::Matrix4f, 2>, 3>>> mat_b  // for J=3, K=2. 现在的问题是如何有效地将它传递给python？此外，我想对多个输入x = [x_0, x_1, ..., x_N] 执行这些计算，结果mat_a(N, I, 4, 4)超出mat_b(N, J, K, 4, 4)预期。每个计算都是独立的，但我认为用 C++x_i重写这个循环可能会更快。x_i另一方面，如果我们在 C++ 中只有固定大小的数组，任务会变得更容易，这个循环也可以转移到 python。这是我的问题的一些虚拟代码（I=5，J=3，K=2）：// example.cpp#include <pybind11/pybind11.h>#include <pybind11/eigen.h>#include <pybind11/stl.h>#include <pybind11/functional.h>#include <pybind11/stl_bind.h>#include <array>#include <vector>#include <Eigen/Dense>Eigen::Matrix4f get_dummy(){    Eigen::Matrix4f mat_a;    mat_a << 1, 2, 3, 4,             5, 6, 7, 8,             9, 8, 7, 6,             5, 4, 3, 2;    return mat_a;}std::pair< std::vector<std::array<Eigen::Matrix4f, 5> >,           std::vector<std::array<std::array<Eigen::Matrix4f, 2>, 3> > >  get_matrices(std::vector<float> & x){    std::vector<std::array<Eigen::Matrix4f, 5> > mat_a(x.size());    std::vector< std::array< std::array< Eigen::Matrix4f, 2>, 3> > mat_b(x.size());    //    for (u_int i=0; i< x.size(); i++)    //        do_stuff(x[i], mat_a[i], mat_b[i]);    mat_a[0][0] = get_dummy();    return std::make_pair(mat_a, mat_b);    }PYBIND11_MODULE(example, m) {    m.def("get_dummy", &get_dummy, pybind11::return_value_policy::reference_internal);    m.def("get_matrices", &get_matrices, pybind11::return_value_policy::reference_internal);}我通过以下方式编译代码：c++ -O3 -Wall -shared -std=c++14 -fPIC `python3 -m pybind11 --includes` example.cpp -o example`python3-config --extension-suffix`

Answer 1

您最好的选择可能是在 python 端创建数据，以便对其进行重新计数和垃圾收集。

test.py

import example

import numpy as np

array = np.zeros((3, 2, 4, 4), 'f4')

example.do_math(array, 3, 2)

print(array[0, 0])

例子.cpp

#define PY_SSIZE_T_CLEAN

#include <Python.h>

#include <Eigen/Dense>

Eigen::Matrix4f get_dummy() {

    Eigen::Matrix4f mat_a;

    mat_a << 1, 2, 3, 4,

             5, 6, 7, 8,

             9, 8, 7, 6,

             5, 4, 3, 2;

    return mat_a;

}

PyObject * example_meth_do_math(PyObject * self, PyObject * args, PyObject * kwargs) {

    static char * keywords[] = {"array", "rows", "cols", NULL};

    PyObject * array;

    int rows, cols;

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "Oii", keywords, &array, &rows, &cols)) {

        return NULL;

    }

    Py_buffer view = {};

    if (PyObject_GetBuffer(array, &view, PyBUF_SIMPLE)) {

        return NULL;

    }

    Eigen::Matrix4f * ptr = (Eigen::Matrix4f *)view.buf;

    for (int i = 0; i < rows; ++i) {

        for (int j = 0; j < cols; ++j) {

            ptr[i * cols + j] = get_dummy();

        }

    }

    PyBuffer_Release(&view);

    Py_RETURN_NONE;

}

PyMethodDef module_methods[] = {

    {"do_math", (PyCFunction)example_meth_do_math, METH_VARARGS | METH_KEYWORDS, NULL},

    {},

};

PyModuleDef module_def = {PyModuleDef_HEAD_INIT, "example", NULL, -1, module_methods};

extern "C" PyObject * PyInit_example() {

    PyObject * module = PyModule_Create(&module_def);

    return module;

}

setup.py

from setuptools import Extension, setup

ext = Extension(

    name='example',

    sources=['./example.cpp'],

    extra_compile_args=['-fpermissive'],

    include_dirs=['.'], # add the path of Eigen

    library_dirs=[],

    libraries=[],

)

setup(

    name='example',

    version='0.1.0',

    ext_modules=[ext],

)

从这里添加第二个参数并使用两个数组进行计算应该是微不足道的。

您可以使用python setup.py develop.

如果你想分发它，你可以创建一个 wheel 文件python setup.py bdist_wheel。

我曾经numpy创建数据，这确保了数据的底层内存是 C 连续的。

这个例子很简单，它使用一个 Matrix4f 指针来迭代一个 3x2 矩阵数组。随意将 转换ptr为Eigen::Array<Eigen::Matrix4f>, 3, 2>。您不能将其强制转换为 an std::vector，因为 an 的内部数据std::vector包含指针。

请注意，std::vector<std::array<...>>内存中没有单个连续数组。改用Eigen::Array。

编辑：

这是一个使用Eigen Array Map:

PyObject * example_meth_do_math(PyObject * self, PyObject * args, PyObject * kwargs) {

    static char * keywords[] = {"array", NULL};

    PyObject * array;

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O", keywords, &array)) {

        return NULL;

    }

    Py_buffer view = {};

    if (PyObject_GetBuffer(array, &view, PyBUF_SIMPLE)) {

        return NULL;

    }

    Eigen::Map<Eigen::Array<Eigen::Matrix4f, 2, 3>> array_map((Eigen::Matrix4f *)view.buf, 2, 3);

    for (int i = 0; i < 2; ++i) {

        for (int j = 0; j < 3; ++j) {

            array_map(i, j) = get_dummy();

        }

    }

    PyBuffer_Release(&view);

    Py_RETURN_NONE;

}

Answer 2

线性代数不会那么流畅（在那里很难击败 Eigen），但会类似于您在 numpy 中所做的（np.dot(A,B)例如。

如果您想坚持使用 Eigen，请注意使用 STL 有一些技术细节。由于您std::array不再能够包含固定数量的矩阵，因此当您移动到std::vector您会遇到对齐问题（诚然，我不完全理解）。很快就会为您提供 xtensor 的有效实现。