将 Eigen::Matrix 的数组从 C++ 返回到 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)。每个 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])

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

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

我尝试了一些道路但没有奏效（对我来说，这并不意味着它们通常不起作用；我只是想不通）：

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

Answer 1

如果您不太依赖于 Eigen，另一种可能性是 xtensor ( found here)。我已经使用了他们的 python 绑定，之前给出了一个直接与 python 通信的示例（found here）。这将具有能够处理更大的多维数组的优势。线性代数不会那么光滑（在那里很难击败 Eigen），但会类似于你在 numpy 中所做的（例如np.dot(A,B)。

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

Answer 2

您最好的选择可能是在 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。

编辑：

这是一个使用EigenArrayMap的函数：

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;
}