adaptive_max_pool3d_backward_single_out_frame Class — pytorch Architecture

Source Code

aten/src/ATen/native/AdaptiveMaxPooling3d.cpp lines 203–240

template <typename scalar_t>
void adaptive_max_pool3d_backward_single_out_frame(
          scalar_t *gradInput_p,
          const scalar_t *gradOutput_p,
          const int64_t *ind_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, 0, [&](int64_t start, int64_t end) {
    for (const auto d : c10::irange(start, end)) {
      scalar_t *gradInput_p_d = gradInput_p + d*isizeT*isizeH*isizeW;
      const scalar_t *gradOutput_p_d = gradOutput_p + d*osizeT*osizeH*osizeW;
      const int64_t *ind_p_d = ind_p + d*osizeT*osizeH*osizeW;

      /* calculate max points */
      int64_t ot = 0, oh = 0, ow = 0;
      for(ot = 0; ot < osizeT; ot++)
      {
        for(oh = 0; oh < osizeH; oh++)
        {
          for(ow = 0; ow < osizeW; ow++)
          {
            /* retrieve position of max */
            int64_t maxp = ind_p_d[ot*osizeH*osizeW + oh*osizeW + ow];

            /* update gradient */
            gradInput_p_d[maxp] += gradOutput_p_d[ot*osizeH*osizeW + oh*osizeW + ow];
          }
        }
      }
    }
  });
}