通过 Pybind11 返回 numpy 数组

一些评论(然后是一个有效的实现)。

• pybind11 围绕 Python 类型(如pybind11::object、pybind11::list和在本例中为pybind11::array_t<T>)的C++对象包装器实际上只是围绕底层 Python 对象指针的包装器。 在这方面，已经有共享指针包装器的角色，因此将其包装在unique_ptr中是没有意义的：直接返回py::array_t<T>对象本质上已经只是返回一个美化的指针。
• pybind11::array_t可以直接从数据指针构造，因此您可以跳过py::buffer_info中间步骤，而直接将形状和步幅提供给pybind11::array_t构造函数。 以这种方式构造的 numpy 数组不会拥有自己的数据，它只会引用它(也就是说，numpyowndata标志将设置为 false)。
• 内存所有权可以与 Python 对象的生命周期相关联，但您仍然需要正确执行释放。 Pybind11 提供了一个py::capsule类来帮助你做到这一点。 您要做的是通过将其指定为array_t的base参数，使numpy数组依赖于此胶囊作为其父类。 这将使 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

我建议使用ndarray。一个基本原则是，除非明确请求，否则永远不会复制基础数据(否则你很快就会遇到巨大的效率低下)。下面是一个正在使用的示例，但还有其他功能我没有展示，包括转换为特征数组(ndarray::asEigen(array))，这使得它非常强大。

页眉：

#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

由于接受的答案似乎已经过时，我将举一个如何使用当前 API 实现它的快速示例。要记住的关键是，如果我们"从头开始"构造一个数组，我们不需要手动分配任何数据，我们只需要为 py：：array_t 构造函数提供我们想要的形状和步幅，它将管理自己的内存。

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
namespace py = pybind11;
PYBIND11_MODULE(numpywrap, m) {
m.def("f", []() {
constexpr size_t elsize = sizeof(double);
size_t shape[3]{100, 1000, 1000};
size_t strides[3]{1000 * 1000 * elsize, 1000 * elsize, elsize};
auto a = py::array_t<double>(shape, strides);
auto view = a.mutable_unchecked<3>();

for(size_t i = 0; i < a.shape(0); i++)
{
for(size_t j = 0; j < a.shape(1); j++)
{
for(size_t k = 0; k < a.shape(2); k++)
{
view(i,j,k) = whatever_data_should_go_here(i,j,k);
}
}
}
return a;
});
}