src/backends/cpu/im2col.hpp - with_cpu - vuron/adept

src/backends/cpu/im2col.hpp
79 строк3 KB
kolkir
AvgPool2d implementation
08 фев 2025, 20:12
08 фев 2025, 20:1261cfbe7
#pragma once

#include <adept/irange.hpp>
#include <adept/param_array.hpp>
#include <adept/threading.hpp>
#include "data_index.hpp"

namespace adept::cpu {

template <typename T>
static void im2col(const T* data_im,
                   const index_t channels,
                   const ParamArray<2>& input_size,
                   const ParamArray<2>& output_size,
                   const ParamArray<2>& kernel_size,
                   const ParamArray<2>& stride_size,
                   const ParamArray<2>& padding_size,
                   const ParamArray<2>& dilation_size,
                   T* data_col,
                   bool is_channels_last = false) {
  const index_t channels_col = channels * kernel_size[0] * kernel_size[1];

  parallel_for<T>(0, channels_col, [&](index_t begin, index_t end) {
    index_t c_im{0}, h_offset{0}, w_offset{0};
    data_index_init(begin, c_im, channels, h_offset, kernel_size[0], w_offset, kernel_size[1]);

    for (auto c : irange(begin, end)) {
      for (auto h : irange(output_size[0])) {
        int64_t h_im = h * stride_size[0] - padding_size[0] + h_offset * dilation_size[0];
        for (auto w : irange(output_size[1])) {
          int64_t w_im = w * stride_size[1] - padding_size[1] + w_offset * dilation_size[1];
          data_col[(c * output_size[0] + h) * output_size[1] + w] =
              (h_im >= 0 && w_im >= 0 && static_cast<index_t>(h_im) < input_size[0] &&
               static_cast<index_t>(w_im) < input_size[1])
                  ? data_im[(c_im * input_size[0] + h_im) * input_size[1] + w_im]
                  : static_cast<T>(0);
        }
      }

      // move to the next index
      data_index_step(c_im, channels, h_offset, kernel_size[0], w_offset, kernel_size[1]);
    }
  });
}

template <typename T>
static void col2im(const T* data_col,
                   const index_t channels,
                   const ParamArray<2>& input_size,
                   const ParamArray<2>& output_size,
                   const ParamArray<2>& kernel_size,
                   const ParamArray<2>& stride_size,
                   const ParamArray<2>& padding_size,
                   const ParamArray<2>& dilation_size,
                   T* data_im,
                   bool is_channels_last = false) {
  std::fill_n(data_im, input_size[0] * input_size[1] * channels, T(0));

  // no prallelization due to the output accumulation
  const index_t channels_col = channels * kernel_size[0] * kernel_size[1];
  for (auto c : irange(channels_col)) {
    index_t w_offset = c % kernel_size[1];
    index_t h_offset = (c / kernel_size[1]) % kernel_size[0];
    index_t c_im = c / kernel_size[0] / kernel_size[1];

    for (auto h : irange(output_size[0])) {
      int64_t h_im = h * stride_size[0] - padding_size[0] + h_offset * dilation_size[0];
      for (auto w : irange(output_size[1])) {
        int64_t w_im = w * stride_size[1] - padding_size[1] + w_offset * dilation_size[1];

        if (h_im >= 0 && static_cast<index_t>(h_im) < input_size[0] && w_im >= 0 &&
            static_cast<index_t>(w_im) < input_size[1])
          data_im[(c_im * input_size[0] + h_im) * input_size[1] + w_im] +=
              data_col[(c * output_size[0] + h) * output_size[1] + w];
      }
    }
  }
}
}  // namespace adept::cpu