fractional_max_pool2d_backward_out_single_batch_frame Class — pytorch Architecture

Source Code

aten/src/ATen/native/FractionalMaxPool2d.cpp lines 222–247

template <typename scalar_t>
void fractional_max_pool2d_backward_out_single_batch_frame(
  scalar_t* gradInput,
  const scalar_t* gradOutput,
  const int64_t* indices,
  int numPlanes,
  int inputW, int inputH,
  int outputW, int outputH) {
  at::parallel_for(0, numPlanes, 0, [&](int64_t start, int64_t end) {
    for (const auto plane : c10::irange(start, end)) {
      scalar_t* gradInputForPlane = gradInput + plane * inputW * inputH;
      const scalar_t* gradOutputForPlane = gradOutput + plane * outputW * outputH;
      const int64_t* indicesForPlane = indices + plane * outputW * outputH;

      for (int h = 0; h < outputH; ++h) {
        for (int w = 0; w < outputW; ++w) {
          int outputIndex = h * outputW + w;
          int64_t index = indicesForPlane[outputIndex];
          AT_ASSERT(index >= 0 && index < static_cast<int64_t>(inputW) * inputH);

          gradInputForPlane[index] += gradOutputForPlane[outputIndex];
        }
      }
    }
  });
}