vuron/adept: Платформа для обучения и вывода нейронных сетей, с API приближенным к PyTorch.

src/nn/maxpool2d.cpp

23 строки733 B

kolkir

MaxPool2d implementation

09 фев 2025, 16:36

09 фев 2025, 16:36b44f8a7

#include <adept/nn/maxpool2d.hpp>
#include "../dispatch/poolNd.hpp"

namespace adept {

MaxPool2dImpl::MaxPool2dImpl(MaxPool2dOptions options, device_t device, dtype_t dtype)
    : options_(options) {
  set_name("MaxPool2d");
}

Variable MaxPool2dImpl::forward(const Variable& input) {
  auto result = max_pool2d_fwd(input.data(), indices_, options_.kernel, options_.stride,
                               options_.padding, options_.dilation);

  Variable var(std::move(result), {input}, "maxpool2d");
  var.set_backward_fn([this, input = input](const auto& out_grad) mutable {
    if (input.requires_grad())
      input.add_grad(max_pool2d_bwd(out_grad, indices_, input.data().shape()));
  });
  return var;
}

}  // namespace adept