Return 从 C++ 到 Python 的 Eigen::Matrix 数组，无需复制

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)。每个 x_i 的计算都是独立的，但我认为在 C++ 中将此循环写在 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`

并且比在 python 中使用它：

import numpy as np
import example

x = np.zeros(1000)

mat_a, mat_b = get_matrices(x)

print(np.shape(mat_a))
print(np.shape(mat_b))
print(mat_a[0][0])

如果我只想 return 一个 Eigen::Matrix 它工作得很快而且据我所知无需复制。但是，当我尝试将 Eigen:Matrices 与 std::array/std::vector 嵌套时，pybind return 是一个嵌套的 numpy 数组列表，而不是一个多维数组。这是预期的，实际上我对它的效果印象深刻，但它对我来说似乎相当慢，尤其是随着数组尺寸的增长。

问题是如何改进它以获得多维 numpy 数组而无需进行不必要的复制。

有些路我试过但没有用（对我来说，这并不意味着它们一般都行不通；我只是想不通）：

使用Eigen::Tensor代替Eigen:Matrix
在 python 中创建矩阵并通过引用将其传递给 C++
为 array, J> 构建自定义包装器

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])

example.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 构建它。

如果你想分发它，你可以用 python setup.py bdist_wheel.

创建一个 wheel 文件

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

此示例保持简单，它使用 Matrix4f 指针迭代 3x2 矩阵数组。随意将 ptr 转换为 Eigen::Array<Eigen::Matrix4f>, 3, 2>。您不能将它转换为 std::vector，因为 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，另一种可能性是 xtensor ( found here). I've used their python bindings before which give an example of communicating directly with python(found here)。这将具有能够处理更大的 multi-dimensional 数组的优势。线性代数不会那么巧妙（在那里很难击败 Eigen），但会类似于您在 numpy 中所做的事情（例如 np.dot(A,B)。

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

Return 从 C++ 到 Python 的 Eigen::Matrix 数组，无需复制

Return Array of Eigen::Matrix from C++ to Python without copying

c++

python

numpy

eigen

pybind11