web-dev-qa-db-fra.com

retour de tableaux numpy via pybind11

J'ai une fonction C++ calculant un grand tenseur que je voudrais retourner à Python comme un tableau NumPy via pybind11 .

D'après la documentation de pybind11, il semble que l'utilisation de STL unique_ptr est souhaitable. Dans l'exemple suivant, la version commentée fonctionne, tandis que celle donnée compile mais échoue au moment de l'exécution ("Impossible de convertir la valeur de retour de la fonction en un Python!")).

Pourquoi la version du smartpointer échoue-t-elle? Quelle est la manière canonique de créer et de renvoyer un tableau NumPy?

PS: En raison de la structure du programme et de la taille du tableau, il est souhaitable de ne pas copier la mémoire mais de créer le tableau à partir d'un pointeur donné. La propriété de la mémoire doit être prise par 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

Quelques commentaires (puis une implémentation fonctionnelle).

  • les enveloppes d'objets C++ de pybind11 autour des types Python (comme pybind11::object, pybind11::list, et, dans ce cas, pybind11::array_t<T>) ne sont en fait que des wrappers autour d'un pointeur d'objet sous-jacent Python. À cet égard, il joue déjà le rôle d'un wrapper de pointeur partagé, et il est donc inutile d'envelopper cela dans un unique_ptr: retour du py::array_t<T> objet directement est déjà essentiellement en train de renvoyer un pointeur glorifié.
  • pybind11::array_t peut être construit directement à partir d'un pointeur de données, vous pouvez donc ignorer py::buffer_info étape intermédiaire et juste donner la forme et les foulées directement au pybind11::array_t constructeur. Un tableau numpy construit de cette façon ne possédera pas ses propres données, il se contentera de les référencer (c'est-à-dire que l'indicateur numpy owndata sera défini sur false).
  • La propriété de la mémoire peut être liée à la durée de vie d'un objet Python, mais vous êtes toujours prêt à faire correctement la désallocation. Pybind11 fournit un py::capsule classe pour vous aider à faire exactement cela. Ce que vous voulez faire est de faire dépendre le tableau numpy de cette capsule comme classe parente en le spécifiant comme l'argument base à array_t. Cela fera que le tableau numpy le référencera, le gardant vivant tant que le tableau lui-même sera vivant, et invoquera la fonction de nettoyage quand il ne sera plus référencé.
  • Le c_style flag dans les versions plus anciennes (antérieures à la version 2.2) n'a eu d'effet que sur les nouveaux tableaux, c'est-à-dire en ne passant pas de pointeur de valeur. Cela a été corrigé dans la version 2.2 pour affecter également les foulées automatiques si vous spécifiez uniquement des formes mais pas des foulées. Cela n'a aucun effet si vous spécifiez les foulées directement vous-même (comme je le fais dans l'exemple ci-dessous).

Donc, en rassemblant les pièces, ce code est un module pybind11 complet qui montre comment vous pouvez accomplir ce que vous recherchez (et inclut une sortie C++ pour démontrer que cela fonctionne effectivement correctement):

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

Compiler cela et l'invoquer depuis Python montre que cela fonctionne:

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