| 295 | |
| 296 | template <typename ctype, typename IdxGetter> |
| 297 | void pooling_backward_avg_impl( |
| 298 | const ctype* __restrict /* src */, const ctype* __restrict /* dst */, |
| 299 | const ctype* __restrict diff, ctype* __restrict grad, size_t N, size_t C, |
| 300 | size_t IH, size_t IW, size_t OH, size_t OW, size_t PH, size_t PW, size_t SH, |
| 301 | size_t SW, size_t FH, size_t FW, bool is_include = true) { |
| 302 | std::memset(grad, 0, sizeof(ctype) * (N * C * IH * IW)); |
| 303 | rep(n, N) rep(c, C) rep(oh, OH) rep(ow, OW) { |
| 304 | size_t count = 0u; |
| 305 | rep(fh, FH) rep(fw, FW) { |
| 306 | size_t ih = -PH + oh * SH + fh; |
| 307 | size_t iw = -PW + ow * SW + fw; |
| 308 | if (ih < IH && iw < IW) |
| 309 | ++count; |
| 310 | } |
| 311 | if (is_include) |
| 312 | count = FH * FW; |
| 313 | if (count == 0u) { |
| 314 | megdnn_throw("The pooling window lies outside completely"); |
| 315 | } |
| 316 | rep(fh, FH) rep(fw, FW) { |
| 317 | size_t ih = -PH + oh * SH + fh; |
| 318 | size_t iw = -PW + ow * SW + fw; |
| 319 | if (ih < IH && iw < IW) { |
| 320 | size_t gi = IdxGetter::get_idx(n, c, ih, iw, N, C, IH, IW); |
| 321 | size_t di = IdxGetter::get_idx(n, c, oh, ow, N, C, OH, OW); |
| 322 | auto& gval = grad[gi]; |
| 323 | auto dval = diff[di]; |
| 324 | gval += dval / ctype(count); |
| 325 | } |
| 326 | } |
| 327 | } |
| 328 | } |
| 329 | |
| 330 | template <typename ctype, typename IdxGetter> |
| 331 | void pooling_backward_avg_expd_impl( |