通过pybind11返回numpy数组

时间:2017-06-20 17:52:25

标签: python c++ arrays numpy pybind11

我有一个计算大张量的C ++函数,我希望通过pybind11将其作为NumPy数组返回到Python。

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

为什么smartpointer版本失败了?创建和返回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) );
}

1 个答案:

答案 0 :(得分:32)

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

  • pybind11围绕Python类型的C ++对象包装器(如pybind11::objectpybind11::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类来帮助您完成此操作。你想要做的是通过将numpy数组指定为base的{​​{1}}参数,使numpy数组依赖于此数据包作为其父类。这将使numpy数组引用它,只要数组本身处于活动状态就保持活动状态,并在不再引用它时调用清理函数。
  • 旧版(2.2之前版本)中的array_t标志仅对新数组产生影响,即未传递值指针时。如果您仅指定形状但不指定步幅,则在2.2版本中修复此选项也会影响自动步幅。如果您自己直接指定步幅,则完全没有效果(正如我在下面的示例中所做的那样)。

因此,将这些部分组合在一起,这段代码是一个完整的pybind11模块,它演示了如何完成您所需要的内容(并包含一些C ++输出以证明其确实正常工作):

c_style

编译并从Python调用它表明它正常工作:

#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();
}