通过pybind11返回numpy数组 [英] returning numpy arrays via pybind11

问题描述

我有一个计算大张量的C ++函数，我想通过 pybind11 .

I have a C++ function computing a large tensor which I would like to return to Python as a NumPy array via pybind11.

从pybind11的文档中，似乎最好使用 STL unique_ptr . 在以下示例中，注释掉的版本有效，而给定的版本可以编译但在运行时失败(无法将函数返回值转换为Python类型！").

From the documentation of pybind11, it seems like using STL unique_ptr is desirable. In the following example, the commented out version works, whereas the given one compiles but fails at runtime ("Unable to convert function return value to a Python type!").

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

Why is the smartpointer version failing? What is the canonical way to create and return a NumPy array?

PS:由于程序的结构和数组的大小，希望不复制内存，而是根据给定的指针创建数组.内存所有权应由Python承担.

PS: Due to program structure and size of the array, it is desirable to not copy memory but create the array from a given pointer. Memory ownership should be taken by 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) );
}

推荐答案

一些注释(然后是有效的实现).

A few comments (then a working implementation).

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类来帮助您准确地做到这一点.您想要做的是通过将其指定为array_t的base参数，使numpy数组依赖于此胶囊作为其父类.这将使numpy数组对其进行引用，并在数组本身处于活动状态时保持其活动状态，并在不再引用该数组时调用清除函数.
• 在较早的版本(2.2之前)中的c_style标志仅对新数组有效，即未传递值指针时.如果仅指定形状而不指定步幅，则该问题在2.2版本中已修复，也将影响自动步幅.如果您直接指定步幅(如我在下面的示例中所做的那样)，则根本没有效果.

• pybind11's C++ object wrappers around Python types (like pybind11::object, pybind11::list, and, in this case, pybind11::array_t<T>) are really just wrappers around an underlying Python object pointer. In this respect there are already taking on the role of a shared pointer wrapper, and so there's no point in wrapping that in a unique_ptr: returning the py::array_t<T> object directly is already essentially just returning a glorified pointer.
• pybind11::array_t can be constructed directly from a data pointer, so you can skip the py::buffer_info intermediate step and just give the shape and strides directly to the pybind11::array_t constructor. A numpy array constructed this way won't own its own data, it'll just reference it (that is, the numpy owndata flag will be set to false).
• Memory ownership can be tied to the life of a Python object, but you're still on the hook for doing the deallocation properly. Pybind11 provides a py::capsule class to help you do exactly this. What you want to do is make the numpy array depend on this capsule as its parent class by specifying it as the base argument to array_t. That will make the numpy array reference it, keeping it alive as long as the array itself is alive, and invoke the cleanup function when it is no longer referenced.
• The c_style flag in the older (pre-2.2) releases only had an effect on new arrays, i.e. when not passing a value pointer. That was fixed in the 2.2 release to also affect the automatic strides if you specify only shapes but not strides. It has no effect at all if you specify the strides directly yourself (as I do in the example below).

因此，将各个部分放在一起，该代码是一个完整的pybind11模块，该模块演示了如何实现所需的功能(并包括一些C ++输出以演示确实可以正常工作):

So, putting the pieces together, this code is a complete pybind11 module that demonstrates how you can accomplish what you're looking for (and includes some C++ output to demonstrate that is indeed working correctly):

#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调用它可以正常工作:

Compiling that and invoking it from Python shows it working:

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

这篇关于通过pybind11返回numpy数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！