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

python/cpp_bindings/conv2d_py.cpp
70 строк3 KB
kolkir
Add python bindings for device configuration
19 июн 2025, 00:03
19 июн 2025, 00:03c00b85c
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/nn/conv2d.hpp>

#include <sstream>

namespace py = pybind11;

using namespace adept;

namespace {
std::shared_ptr<Conv2dImpl> make_conv2d(index_t in_ch,
                                        index_t out_ch,
                                        const std::vector<index_t>& kernel,
                                        const std::vector<index_t>& stride,
                                        const std::vector<index_t>& padding,
                                        const std::vector<index_t>& dilation,
                                        bool bias,
                                        device_t device,
                                        dtype_t dtype,
                                        index_t device_id) {
  return std::make_shared<Conv2dImpl>(Conv2dOptions(in_ch, out_ch, kernel)
                                          .with_stride(stride)
                                          .with_padding(padding)
                                          .with_dilation(dilation)
                                          .with_bias(bias),
                                      device, dtype, device_id);
}
std::shared_ptr<Conv2dImpl> make_conv2d_sv(index_t in_ch,
                                           index_t out_ch,
                                           index_t kernel,
                                           index_t stride,
                                           index_t padding,
                                           index_t dilation,
                                           bool bias,
                                           device_t device,
                                           dtype_t dtype,
                                           index_t device_id) {
  return std::make_shared<Conv2dImpl>(Conv2dOptions(in_ch, out_ch, kernel)
                                          .with_stride(stride)
                                          .with_padding(padding)
                                          .with_dilation(dilation)
                                          .with_bias(bias),
                                      device, dtype, device_id);
}
}  // namespace

void bind_conv2d(py::module& m) {
  py::class_<Conv2dImpl, Module, std::shared_ptr<Conv2dImpl>>(m, "Conv2d")
      .def(py::init(&make_conv2d), py::arg("in_channels"), py::arg("out_channels"),
           py::arg("kernel"), py::arg("stride") = std::vector<index_t>{1, 1},
           py::arg("padding") = std::vector<index_t>{0, 0},
           py::arg("dilation") = std::vector<index_t>{1, 1}, py::arg("bias") = true,
           py::arg("device") = device_t::CPU, py::arg("dtype") = dtype_t::Float32,
           py::arg("device_id") = 0)
      .def(py::init(&make_conv2d_sv), py::arg("in_channels"), py::arg("out_channels"),
           py::arg("kernel"), py::arg("stride") = index_t{1}, py::arg("padding") = index_t{0},
           py::arg("dilation") = index_t{1}, py::arg("bias") = true,
           py::arg("device") = device_t::CPU, py::arg("dtype") = dtype_t::Float32,
           py::arg("device_id") = 0)
      .def("parameters", &Conv2dImpl::parameters)
      .def("forward", &Conv2dImpl::forward)
      .def("__call__", [](Conv2dImpl& self, Variable& input) { return self.forward(input); })
      .def("__repr__", [](const Conv2dImpl& v) {
        std::stringstream buf;
        buf << v;
        return buf.str();
      });
}