adaptive_avg_pool3d_backward_out_frame Class — pytorch Architecture

Source Code

aten/src/ATen/native/AdaptiveAveragePooling3d.cpp lines 178–228

template <typename scalar_t>
void adaptive_avg_pool3d_backward_out_frame(
    scalar_t* gradInput_p,
    const scalar_t* gradOutput_p,
    int64_t sizeD,
    int64_t isizeT,
    int64_t isizeH,
    int64_t isizeW,
    int64_t osizeT,
    int64_t osizeH,
    int64_t osizeW) {
  at::parallel_for(0, sizeD, 1, [&](int64_t start, int64_t end) {
    for (const auto d : c10::irange(start, end)) {
      scalar_t* gradInput_p_d = gradInput_p + d * isizeT * isizeW * isizeH;
      const scalar_t* gradOutput_p_d = gradOutput_p + d * osizeT * osizeW * osizeH;

      /* calculate average */
      for (const auto ot : c10::irange(osizeT)) {
        auto istartT = start_index(ot, osizeT, isizeT);
        auto iendT = end_index(ot, osizeT, isizeT);
        auto kT = iendT - istartT;

        for (const auto oh : c10::irange(osizeH)) {
          auto istartH = start_index(oh, osizeH, isizeH);
          auto iendH = end_index(oh, osizeH, isizeH);
          auto kH = iendH - istartH;

          for (const auto ow : c10::irange(osizeW)) {
            auto istartW = start_index(ow, osizeW, isizeW);
            auto iendW = end_index(ow, osizeW, isizeW);
            auto kW = iendW - istartW;

            scalar_t grad_delta =
                gradOutput_p_d[ot * osizeH * osizeW + oh * osizeW + ow] / kT /
                kH / kW;

            for (const auto it : c10::irange(istartT, iendT)) {
              for (const auto ih : c10::irange(istartH, iendH)) {
                for (const auto iw : c10::irange(istartW, iendW)) {
                  /* update gradient */
                  gradInput_p_d[it * isizeH * isizeW + ih * isizeW + iw] +=
                      grad_delta;
                }
              }
            }
          }
        }
      }
    }
  });
}