通过 pybind11 返回 numpy 数组

Question

我有一个 C++ 函数计算一个大张量，我想通过 pybind11 return 到 Python 作为 NumPy 数组。

从 pybind11 的文档来看，使用 STL unique_ptr 似乎是可取的。在以下示例中，注释掉的版本有效，而给定版本编译但在运行时失败 ("Unable to convert function return value to a Python type!")。

为什么智能指针版本失败？创建和 return NumPy 数组的规范方法是什么？

PS：由于程序结构和数组的大小，最好不要复制内存，而是从给定的指针创建数组。内存所有权应由 Python.

取得

typedef typename py::array_t<double, py::array::c_style | py::array::forcecast> py_cdarray_t;

// py_cd_array_t _test()
std::unique_ptr<py_cdarray_t> _test()
{
    double * memory = new double[3]; memory[0] = 11; memory[1] = 12; memory[2] = 13;
    py::buffer_info bufinfo (
        memory,                                   // pointer to memory buffer
        sizeof(double),                           // size of underlying scalar type
        py::format_descriptor<double>::format(),  // python struct-style format descriptor
        1,                                        // number of dimensions
        { 3 },                                    // buffer dimensions
        { sizeof(double) }                        // strides (in bytes) for each index
    );

    //return py_cdarray_t(bufinfo);
    return std::unique_ptr<py_cdarray_t>( new py_cdarray_t(bufinfo) );
}

Answer 1

一些评论（然后是一个有效的实现）。

pybind11 的 C++ 对象包装 Python 类型（如 pybind11::object、pybind11::list，在这种情况下，pybind11::array_t<T>）实际上只是对底层 Python 对象指针。在这方面，已经承担了共享指针包装器的角色，因此将其包装在 unique_ptr 中没有意义：直接返回 py::array_t<T> 对象本质上已经只是返回一个美化的指针。
pybind11::array_t 可以直接从数据指针构造，因此您可以跳过 py::buffer_info 中间步骤，直接将形状和步长提供给 pybind11::array_t 构造函数。以这种方式构造的 numpy 数组不会拥有自己的数据，它只会引用它（也就是说，numpy owndata 标志将设置为 false）。
内存所有权可以与 Python 对象的生命周期相关联，但是您仍然需要正确地进行释放。 Pybind11 提供了一个 py::capsule class 来帮助你做到这一点。您想要做的是通过将它指定为 array_t 的 base 参数，使 numpy 数组依赖于此胶囊作为其父级 class。这将使 numpy 数组引用它，只要数组本身还活着，它就会保持活动状态，并在不再被引用时调用清理函数。
旧版本（2.2 之前）中的 c_style 标志仅对新数组有影响，即不传递值指针时。这在 2.2 版本中已修复，如果您仅指定形状而不指定步幅，也会影响自动步幅。如果你自己直接指定步幅（就像我在下面的例子中所做的那样），它根本没有效果。

因此，将各个部分放在一起，此代码是一个完整的 pybind11 模块，它演示了如何完成您正在寻找的内容（并包含一些 C++ 输出以证明它确实可以正常工作）：

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

namespace py = pybind11;

PYBIND11_PLUGIN(numpywrap) {
    py::module m("numpywrap");
    m.def("f", []() {
        // Allocate and initialize some data; make this big so
        // we can see the impact on the process memory use:
        constexpr size_t size = 100*1000*1000;
        double *foo = new double[size];
        for (size_t i = 0; i < size; i++) {
            foo[i] = (double) i;
        }

        // Create a Python object that will free the allocated
        // memory when destroyed:
        py::capsule free_when_done(foo, [](void *f) {
            double *foo = reinterpret_cast<double *>(f);
            std::cerr << "Element [0] = " << foo[0] << "\n";
            std::cerr << "freeing memory @ " << f << "\n";
            delete[] foo;
        });

        return py::array_t<double>(
            {100, 1000, 1000}, // shape
            {1000*1000*8, 1000*8, 8}, // C-style contiguous strides for double
            foo, // the data pointer
            free_when_done); // numpy array references this parent
    });
    return m.ptr();
}

编译它并从 Python 调用它显示它工作：

>>> import numpywrap
>>> z = numpywrap.f()
>>> # the python process is now taking up a bit more than 800MB memory
>>> z[1,1,1]
1001001.0
>>> z[0,0,100]
100.0
>>> z[99,999,999]
99999999.0
>>> z[0,0,0] = 3.141592
>>> del z
Element [0] = 3.14159
freeing memory @ 0x7fd769f12010
>>> # python process memory size has dropped back down

Answer 2

我建议使用 ndarray。一个基本原则是，除非明确要求，否则永远不会复制底层数据（否则您很快就会陷入效率低下的境地）。下面是它的一个使用示例，但还有其他我没有展示的功能，包括转换为 Eigen 数组 (ndarray::asEigen(array))，这使得它非常强大。

Header:

#ifndef MYTENSORCODE_H
#define MYTENSORCODE_H

#include "ndarray_fwd.h"

namespace myTensorNamespace {

ndarray::Array<double, 2, 1> myTensorFunction(int param1, double param2);

}  // namespace myTensorNamespace

#endif  // include guard

图书馆：

#include "ndarray.h"
#include "myTensorCode.h"

namespace myTensorNamespace {

ndarray::Array<double, 2, 1> myTensorFunction(int param1, double param2) {
    std::size_t const size = calculateSize();
    ndarray::Array<double, 2, 1> array = ndarray::allocate(size, size);
    array.deep() = 0;  // initialise
    for (std::size_t ii = 0; ii < size; ++ii) {
        array[ii][ndarray::view(ii, ii + 1)] = 1.0;
    }
    return array;
}

}  // namespace myTensorNamespace

包装器：

#include "pybind11/pybind11.h"
#include "ndarray.h"
#include "ndarray/pybind11.h"

#include "myTensorCode.h"

namespace py = pybind11;
using namespace pybind11::literals;

namespace myTensorNamespace {
namespace {

PYBIND11_MODULE(myTensorModule, mod) {
    mod.def("myTensorFunction", &myTensorFunction, "param1"_a, "param2"_a);
}

}  // anonymous namespace
}  // namespace myTensorNamespace

通过 pybind11 返回 numpy 数组

returning numpy arrays via pybind11

c++

python

arrays

numpy

pybind11