cpu_avg_pool2d Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/AvgPoolKernel.cpp lines 16–100

template <typename scalar_t>
void cpu_avg_pool2d(
    const Tensor& output_,
    const Tensor& input_,
    int64_t kW, int64_t kH,
    int64_t dW, int64_t dH,
    int64_t padW, int64_t padH,
    bool count_include_pad,
    std::optional<int64_t> divisor_override) {
  using acc_t = at::opmath_type<scalar_t>;

  auto input = input_.contiguous();
  auto output = output_.contiguous();

  auto input_data = input.const_data_ptr<scalar_t>();
  auto output_data = output.data_ptr<scalar_t>();

  int64_t numel = output.numel();
  int64_t ndim = input.ndimension();
  // treat batch size and channels as one dimension
  int64_t channels = ndim == 3 ? input.size(0) : input.size(0) * input.size(1);
  int64_t input_height = input.size(-2);
  int64_t input_width = input.size(-1);
  int64_t output_height = output.size(-2);
  int64_t output_width = output.size(-1);

  // parallel on dim N, C, H, W
  at::parallel_for(0, numel, 0, [&](int64_t begin, int64_t end) {
    int64_t c = 0;
    int64_t oh = 0;
    int64_t ow = 0;
    data_index_init(begin, c, channels, oh, output_height, ow, output_width);

    for (const auto i : c10::irange(begin, end)) {
      output_data[i] = static_cast<scalar_t>(0);

      // local pointers
      const scalar_t* input_ptr = input_data + c * input_height * input_width;

      // compute the mean of the input image...
      int64_t ih0 = oh * dH - padH;
      int64_t iw0 = ow * dW - padW;
      int64_t ih1 = std::min(ih0 + kH, input_height + padH);
      int64_t iw1 = std::min(iw0 + kW, input_width + padW);
      int64_t pool_size = (ih1 - ih0) * (iw1 - iw0);
      ih0 = std::max(ih0, (int64_t) 0);
      iw0 = std::max(iw0, (int64_t) 0);
      ih1 = std::min(ih1, input_height);
      iw1 = std::min(iw1, input_width);

      if (ih0 >= ih1 || iw0 >= iw1) {
        // move on to next output index
        data_index_step(c, channels, oh, output_height, ow, output_width);
        continue;
      }

      acc_t sum = 0;

      int64_t divide_factor = 0;
      if (divisor_override.has_value()) {
        divide_factor = divisor_override.value();
      } else {
        if(count_include_pad) {
          divide_factor = pool_size;
        } else {
          divide_factor = (ih1 - ih0) * (iw1 - iw0);
        }
      }

      for (const auto ih : c10::irange(ih0, ih1)) {
        for (const auto iw : c10::irange(iw0, iw1)) {
          sum += input_ptr[ih * input_width + iw];
        }
      }
      output_data[i] += scalar_t(sum / divide_factor);

      // move on to next output index
      data_index_step(c, channels, oh, output_height, ow, output_width);
    }
  });

  if (!output_.is_contiguous()) {
    output_.copy_(output);
  }
}