src/nn/conv2d.cpp - master - vuron/adept

src/nn/conv2d.cpp
85 строк2 KB
kolkir
revert compiler compatibility to gcc12
03 мар 2025, 23:52
03 мар 2025, 23:5226a9274
#include <adept/nn/conv2d.hpp>
#include <adept/nn/init.hpp>
#include "../dispatch/convNd.hpp"

#include <cmath>

namespace adept {

Conv2dImpl::Conv2dImpl(Conv2dOptions options, device_t device, dtype_t dtype)
    : options_(options),
      weight_(Tensor::empty({.shape = {options_.out_channels, options_.in_channels,
                                       options_.kernel[0], options_.kernel[1]},
                             .device = device,
                             .dtype = dtype})) {
  if (options_.bias) {
    bias_ = Variable(
        Tensor::empty({.shape = {1, options_.out_channels}, .device = device, .dtype = dtype}));
  }
  init_weights();
  set_name("Conv2d");
  register_parameter("weight", weight_);
  register_parameter("bias", bias_);
}

void Conv2dImpl::set_weights(const Tensor& weight) {
  weight_.data() = weight;
}

void Conv2dImpl::set_bias(const Tensor& bias) {
  bias_.data() = bias;
}

Variable Conv2dImpl::weights() const {
  return weight_;
}

Variable Conv2dImpl::bias() const {
  return bias_;
}

void Conv2dImpl::init_weights() {
  CHECK(options_.in_channels > 0 && options_.out_channels > 0, "in_channels=", options_.in_channels,
        " and out_channels=", options_.out_channels, " must be a positive integer.");

  auto in_features = weight_.data().properties().shape.dim(1);

  // kamming He
  auto a = std::sqrt(5);
  const auto gain = std::sqrt(2.0 / (1 + pow(a, 2)));
  auto std = gain / std::sqrt(in_features);
  fill_uniform(-std, std, weight_.data());

  if (options_.bias) {
    std = static_cast<float32_t>(1. / std::sqrt(in_features));
    fill_uniform(-std, std, bias_.data());
  }
}

Variable Conv2dImpl::forward(const Variable& input) {
  Tensor bias;
  if (bias_.defined())
    bias = bias_.data();

  auto result = conv2d_fwd(input.data(), weight_.data(), bias, options_.kernel, options_.stride,
                           options_.padding, options_.dilation);

  Variable var(std::move(result), {input}, "conv2d");
  var.set_backward_fn([this, input = input, bias = bias](const auto& out_grad) mutable {
    Tensor input_grad;
    if (input.requires_grad())
      input_grad = input.grad();

    Tensor bias_grad;
    if (bias_.defined()) {
      bias_grad = bias_.grad();
    }

    auto weight_grad = weight_.grad();
    conv2d_bwd(out_grad, input.data(), weight_.data(), input_grad, weight_grad, bias_grad,
               options_.kernel, options_.stride, options_.padding, options_.dilation);
  });
  return var;
}

}  // namespace adept