kSpatialDim Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/AdaptiveAveragePooling.cpp lines 174–244

template <int32_t kSpatialDim, typename scalar_t>
Tensor _adaptive_avg_pool(const Tensor& input,
                          IntArrayRef output_size,
                          Tensor& output) {
  const auto output_shape = get_output_shape<kSpatialDim>(input, output_size);
  /* sizes */
  int64_t sizeC = input.size(-(kSpatialDim + 1));
  int64_t isizeD = kSpatialDim == 2 ? 1 : input.size(-3);
  int64_t isizeH = input.size(-2);
  int64_t isizeW = input.size(-1);

  auto osizeD = kSpatialDim == 2 ? 1 : output_shape[output_shape.size() - 3];
  auto osizeH = output_shape[output_shape.size() - 2];
  auto osizeW = output_shape[output_shape.size() - 1];

  int64_t sizeB = output_shape.size() ==(kSpatialDim + 1) ? 1 : output_shape[0];
  if (input.is_contiguous(c10::MemoryFormat::ChannelsLast) ||
      input.is_contiguous(c10::MemoryFormat::ChannelsLast3d)) {
    // Fast path for NDHWC
    auto in_stride = input.strides();
    output = at::_empty_affine_quantized(
        output_shape,
        input.options().memory_format(input.suggest_memory_format()),
        input.q_scale(),
        input.q_zero_point(),
        std::nullopt);

    qadaptive_avg_pool3d_ndhwc_stub(
        input.device().type(),
        input,
        output,
        sizeB,
        sizeC,
        isizeD,
        isizeH,
        isizeW,
        osizeD,
        osizeH,
        osizeW,
        in_stride[0],
        in_stride[in_stride.size() - (kSpatialDim + 1)],
        in_stride[in_stride.size() - kSpatialDim],
        in_stride[in_stride.size() - 2],
        in_stride[in_stride.size() - 1]);
    return output;
  } else {
    output = at::_empty_affine_quantized(
        output_shape, input.options(), input.q_scale(), input.q_zero_point());
    auto input_contig = input.contiguous();
    auto input_data = input_contig.data_ptr<scalar_t>();
    auto output_data = output.data_ptr<scalar_t>();
    auto in_stride = input_contig.strides();

    adaptive_avg_pool_single_out_frame<scalar_t>(
        input_data,
        output_data,
        // Contract batch and channels into one dimension
        sizeB * sizeC,
        isizeD,
        isizeH,
        isizeW,
        osizeD,
        osizeH,
        osizeW,
        in_stride[in_stride.size() - (kSpatialDim + 1)],
        in_stride[in_stride.size() - kSpatialDim],
        in_stride[in_stride.size() - 2],
        in_stride[in_stride.size() - 1]);
    return output;
  }
}