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

src/nn/avgpool2d.cpp

23 строки751 B

kolkir

AvgPool2d implementation

08 фев 2025, 20:12

08 фев 2025, 20:1261cfbe7

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

namespace adept {

AvgPool2dImpl::AvgPool2dImpl(AvgPool2dOptions options, device_t device, dtype_t dtype)
    : options_(options) {
  set_name("AvgPool2d");
}

Variable AvgPool2dImpl::forward(const Variable& input) {
  auto result = avg_pool2d_fwd(input.data(), options_.kernel, options_.stride, options_.padding);

  Variable var(std::move(result), {input}, "avgpool2d");
  var.set_backward_fn([this, input = input](const auto& out_grad) mutable {
    if (input.requires_grad())
      input.add_grad(avg_pool2d_bwd(out_grad, input.data().shape(), options_.kernel,
                                    options_.stride, options_.padding));
  });
  return var;
}

}  // namespace adept