python — возвращает массив numpy через pybind11

У меня есть функция C ++, вычисляющая большой тензор, который я хотел бы вернуть в Python в виде массива NumPy через pybind11.

Из документации 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) );
}

13

Решение

Несколько комментариев (потом рабочая реализация).

  • Оболочки объектов C ++ в Pybind11 вокруг типов 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 флаг будет установлен в ложь).
  • Владение памятью может быть привязано к жизни объекта Python, но вы все еще на крючке для правильного выполнения освобождения. Pybind11 предоставляет py::capsule класс, чтобы помочь вам сделать именно это. То, что вы хотите сделать, это сделать массив numpy зависимым от этой капсулы как ее родительского класса, указав его как base аргумент array_t, Это будет ссылаться на массив numpy, сохраняя его живым до тех пор, пока активен сам массив, и вызывать функцию очистки, когда на него больше нет ссылок.
  • c_style Флаг в более ранних (до 2.2) выпусках влиял только на новые массивы, т.е. когда не передается указатель значения. Это было исправлено в версии 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
25

Другие решения

Других решений пока нет …

По вопросам рекламы [email protected]