web-dev-qa-db-ja.com

pybind11経由でnumpy配列を返す

pybind11 を介してNumPy配列としてPython)に戻したい大きなテンソルを計算するC++関数があります。

Pybind11のドキュメントから、 STL unique_ptr を使用することが望ましいようです。次の例では、コメントアウトされたバージョンが機能しますが、指定されたバージョンはコンパイルされますが、実行時に失敗します(「関数の戻り値をPython type!)に変換できません!」)。

スマートポインターのバージョンが失敗するのはなぜですか? 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) );
}
22
mrupp

いくつかのコメント(それから実際の実装)。

  • pybind11のPython型(pybind11::objectpybind11::list、この場合はpybind11::array_t<T>など)の周りのC++オブジェクトラッパー)は、実際には基礎となる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配列をこのカプセルにその親クラスとして依存させることです。クリーンアップ関数が参照されなくなったら呼び出します。
  • 古い(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
38