python/cpp_bindings/maxpool2d_py.cpp - with_cpu - vuron/adept

python/cpp_bindings/maxpool2d_py.cpp
45 строк2 KB
kolkir
Initial ResNet implementation
06 мар 2025, 00:15
06 мар 2025, 00:15155bcba
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/nn/maxpool2d.hpp>

#include <sstream>

namespace py = pybind11;

using namespace adept;

namespace {
std::shared_ptr<MaxPool2dImpl> make_maxpool2d(const std::vector<index_t>& kernel,
                                              const std::vector<index_t>& stride,
                                              const std::vector<index_t>& padding,
                                              const std::vector<index_t>& dilation) {
  return std::make_shared<MaxPool2dImpl>(
      MaxPool2dOptions(kernel).with_stride(stride).with_padding(padding).with_dilation(dilation));
}
std::shared_ptr<MaxPool2dImpl> make_maxpool2d_sv(index_t kernel,
                                                 index_t stride,
                                                 index_t padding,
                                                 index_t dilation) {
  return std::make_shared<MaxPool2dImpl>(
      MaxPool2dOptions(kernel).with_stride(stride).with_padding(padding).with_dilation(dilation));
}
}  // namespace

void bind_maxpool2d(py::module& m) {
  py::class_<MaxPool2dImpl, Module, std::shared_ptr<MaxPool2dImpl>>(m, "MaxPool2d")
      .def(py::init(&make_maxpool2d), 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})
      .def(py::init(&make_maxpool2d_sv), py::arg("kernel"), py::arg("stride") = index_t{1},
           py::arg("padding") = index_t{0}, py::arg("dilation") = index_t{1})
      .def("parameters", &MaxPool2dImpl::parameters)
      .def("forward", &MaxPool2dImpl::forward)
      .def("__call__", [](MaxPool2dImpl& self, Variable& input) { return self.forward(input); })
      .def("__repr__", [](const MaxPool2dImpl& v) {
        std::stringstream buf;
        buf << v;
        return buf.str();
      });
}