python/cpp_bindings/module_py.cpp - master - glebestraikh/adept

python/cpp_bindings/module_py.cpp
55 строк2 KB
kolkir
Fix python bindings to support GPU backend
06 июл 2025, 18:37
06 июл 2025, 18:374c42c77
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/nn/module.hpp>

namespace py = pybind11;

using namespace adept;

void bind_module(py::module& m) {
  py::class_<Module, std::shared_ptr<Module>>(m, "Module")
      .def(py::init<std::string>())
      .def("name", &Module::name)
      .def("parameters", &Module::parameters)
      .def("save", [](const Module& self, OutputSerializer& output) { self.save(output); })
      .def("load", [](Module& self, InputSerializer& input) { self.load(input); })
      .def("train", &Module::train)
      .def("eval", &Module::eval)
      .def("cpu", &Module::cpu)
      .def("gpu", &Module::gpu, py::arg("device_id") = 0)
      .def("to", &Module::to, py::arg("device"), py::arg("device_id") = 0)
      .def("register_module",
           [](Module& self, const std::string_view name, py::object value) -> void {
             try {
               auto m = value.cast<std::shared_ptr<Module>>();
               if (m)  // ignore None values
                 self.register_module(std::string(name), m);
             } catch (const py::cast_error&) {
               // ignore
             }
           })
      .def("__repr__",
           [](const Module& v) {
             std::stringstream buf;
             buf << v;
             return buf.str();
           })
      .def("__setattr__",
           [](Module& instance, const std::string_view name, py::object value) -> void {
             try {
               auto m = value.cast<std::shared_ptr<Module>>();
               if (m)  // ignore None values
                 instance.register_module(std::string(name), m);
             } catch (const py::cast_error&) {
               // ignore
             }

             // fall back on parent class's __setattr__ ...
             py::type selfType = py::type::of<Module>();
             py::object selfObj = py::cast(instance);
             py::module::import("builtins")
                 .attr("super")(selfType, selfObj)
                 .attr("__setattr__")(name, value);
           });
}