返回多个 py::array 而不在 pybind11 中复制
Posted
技术标签:
【中文标题】返回多个 py::array 而不在 pybind11 中复制【英文标题】:returning multiple py::array without copying in pybind11 【发布时间】:2019-09-26 09:43:43 【问题描述】:我正在尝试使用 pybind11 在 C++ 中构建一个 python 模块。我有以下代码:
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
namespace py = pybind11;
struct ContainerElement
uint8_t i;
double d;
double d2;
;
class Container
private:
std::vector<uint8_t> ints;
std::vector<double> doubles;
std::vector<double> doubles2;
public:
std::vector<uint8_t>& getInts() return ints;
std::vector<double>& getDoubles() return doubles;
std::vector<double>& getDoubles2() return doubles2;
void addElement(ContainerElement element)
ints.emplace_back(element.i);
doubles.emplace_back(element.d);
doubles2.emplace_back(element.d2);
;
void fillContainer(Container& container)
for (int i = 0; i < 1e6; ++i)
container.addElement((uint8_t)i, (double)i,(double)i );
PYBIND11_MODULE(containerInterface, m)
py::class_<Container>(m, "Container")
.def(py::init<>())
.def("getInts", [](Container& container)
return py::array_t<uint8_t>(
container.getInts().size() ,
sizeof(uint8_t) ,
container.getInts().data());
)
.def("getDoubles", [](Container& container)
return py::array_t<double>(
container.getDoubles().size() ,
sizeof(double) ,
container.getDoubles().data());
)
.def("getDoubles2", [](Container& container)
return py::array_t<double>(
container.getDoubles2().size() ,
sizeof(double) ,
container.getDoubles2().data());
);
m.def("fillContainer", &fillContainer);
当我在 python 中调用这段代码时:
import containerInterface
container = containerInterface.Container()
containerInterface.fillContainer(container)
i = container.getInts()
d = container.getDoubles()
d2 = container.getDoubles2()
这可行,但是当我检查程序的内存使用情况时(使用psutil.Process(os.getpid()).memory_info().rss
),当我调用函数getInts, getDoubles
和getDoubles2
时似乎会进行复制。有没有办法避免这种情况?
我尝试过使用np.array(container.getInts(), copy=False)
,但它仍然会复制。我还尝试在 Container 类上使用py::buffer_protocol()
,如此处所述:https://pybind11.readthedocs.io/en/stable/advanced/pycpp/numpy.html。但是,我只能对 Ints 向量或 Doubles 向量进行此操作,而不能同时对所有向量进行。
PYBIND11_MODULE(containerInterface, m)
py::class_<Container>(m, "Container", py::buffer_protocol())
.def(py::init<>())
.def("getInts", &Container::getInts)
.def("getDoubles", &Container::getDoubles)
.def_buffer([](Container& container) -> py::buffer_info
return py::buffer_info(
container.getInts().data(),
sizeof(uint8_t),
py::format_descriptor<uint8_t>::format(),
1,
container.getInts().size() ,
sizeof(uint8_t) * container.getInts().size()
);
);
m.def("fillContainer", &fillContainer);
然后我可以使用i = np.array(container, copy=False)
,而无需复制。但是,正如我所说,它现在仅适用于 Ints
向量。
【问题讨论】:
【参考方案1】:我找到了一个可行的解决方案。虽然它可能不是最优雅的。我创建了三个新类Ints
、Doubles
和Doubles2
,它们采用原始容器并通过函数调用getValues()
公开各自的向量。使用这三个类,我可以为所有类指定 3 次缓冲区协议。
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <pybind11/buffer_info.h>
namespace py = pybind11;
struct ContainerElement
uint8_t i;
double d;
double d2;
;
class Container
private:
std::vector<uint8_t> ints;
std::vector<double> doubles;
std::vector<double> doubles2;
public:
std::vector<uint8_t>& getInts() return ints;
std::vector<double>& getDoubles() return doubles;
std::vector<double>& getDoubles2() return doubles2;
void addElement(ContainerElement element)
ints.emplace_back(element.i);
doubles.emplace_back(element.d);
doubles2.emplace_back(element.d2);
;
void fillContainer(Container& container)
for (int i = 0; i < 1e6; ++i)
container.addElement( (uint8_t)i, (double)i,(double)i );
class Ints
private:
Container& cont;
public:
Ints(Container& cont) : cont(cont)
std::vector<uint8_t>& getValues() return cont.getInts();
;
class Doubles
private:
Container& cont;
public:
Doubles(Container& cont) : cont(cont)
std::vector<double>& getValues() return cont.getDoubles();
;
class Doubles2
private:
Container& cont;
public:
Doubles2(Container& cont) : cont(cont)
std::vector<double>& getValues() return cont.getDoubles2();
;
PYBIND11_MODULE(newInterface, m)
py::class_<Container>(m, "Container")
.def(py::init<>());
py::class_<Ints>(m, "Ints", py::buffer_protocol())
.def(py::init<Container&>(), py::keep_alive<1, 2>())
.def_buffer([](Ints& ints) -> py::buffer_info
return py::buffer_info(
ints.getValues().data(),
sizeof(uint8_t),
py::format_descriptor<uint8_t>::format(),
ints.getValues().size()
);
);
py::class_<Doubles>(m, "Doubles", py::buffer_protocol())
.def(py::init<Container&>(), py::keep_alive<1, 2>())
.def_buffer([](Doubles& doubles) -> py::buffer_info
return py::buffer_info(
doubles.getValues().data(),
sizeof(double),
py::format_descriptor<double>::format(),
doubles.getValues().size()
);
);
py::class_<Doubles2>(m, "Doubles2", py::buffer_protocol())
.def(py::init<Container&>(), py::keep_alive<1, 2>())
.def_buffer([](Doubles2& doubles2) -> py::buffer_info
return py::buffer_info(
doubles2.getValues().data(),
sizeof(double),
py::format_descriptor<double>::format(),
doubles2.getValues().size()
);
);
m.def("fillContainer", &fillContainer);
这样我就可以在 Python 中按以下方式使用代码:
import newInterface as ci
import numpy as np
container = ci.Container()
ci.fillContainer(container)
i = np.array(ci.Ints(container), copy=False)
d = np.array(ci.Doubles(container), copy=False)
d2 = np.array(ci.Doubles2(container), copy=False)
一旦fillContainer
填充了容器,numpy 数组的构造就不会从这个容器中复制值。
【讨论】:
【参考方案2】:我猜您必须指定访问函数返回引用而不是副本,这可能是默认值。我不知道你是如何用 pybind 做到这一点的,但我已经用 boost::python 和 Ponder 做到了这一点。
即您需要指定退货政策。
【讨论】:
您可能正在做某事,但是将返回值政策更改为例如return_value_policy::move
似乎没有任何改变。【参考方案3】:
这并不能直接解决问题,但仍然允许在不进行复制的情况下返回数组缓冲区。 灵感来自这个线程: https://github.com/pybind/pybind11/issues/1042
基本上,只需向 py::array() 构造函数提供一个 py::capsule。 有了这个,py::array() 构造函数不分配一个单独的缓冲区和副本。例如:
// Use this if the C++ buffer should NOT be deallocated
// once Python no longer has a reference to it
py::capsule buffer_handle([]());
// Use this if the C++ buffer SHOULD be deallocated
// once the Python no longer has a reference to it
// py::capsule buffer_handle(data_buffer, [](void* p) free(p); );
return py::array(py::buffer_info(
data_buffer,
element_size,
data_type,
dims_length,
dims,
strides
), buffer_handle);
【讨论】:
以上是关于返回多个 py::array 而不在 pybind11 中复制的主要内容,如果未能解决你的问题,请参考以下文章
Pybind11:从 C++ 端创建并返回 numpy 数组