| 34 | |
| 35 | template <typename T> |
| 36 | void backward( |
| 37 | _megdnn_tensor_in src, _megdnn_tensor_in /* dst */, _megdnn_tensor_in diff, |
| 38 | _megdnn_tensor_out grad, const Param& param) { |
| 39 | auto N = src.layout.shape[0], C = src.layout.shape[1], H = src.layout.shape[2], |
| 40 | W = src.layout.shape[3]; |
| 41 | auto half_window = param.n / 2; |
| 42 | auto k = param.k, alpha = param.alpha, beta = param.beta; |
| 43 | auto sptr = src.ptr<T>(), hptr = diff.ptr<T>(), gptr = grad.ptr<T>(); |
| 44 | std::fill(gptr, gptr + N * C * H * W, 0); |
| 45 | rep(n, N) rep(hw, H * W) { |
| 46 | rep(dc, C) { |
| 47 | auto didx = n * C * H * W + dc * H * W + hw; |
| 48 | size_t sc_start = (dc >= half_window ? dc - half_window : 0u); |
| 49 | size_t sc_end = std::min(dc + half_window, C - 1); |
| 50 | float tmp = k; |
| 51 | for (size_t sc = sc_start; sc <= sc_end; ++sc) { |
| 52 | auto sidx = n * C * H * W + sc * H * W + hw; |
| 53 | tmp += alpha * sqr(sptr[sidx]); |
| 54 | } |
| 55 | for (size_t sc = sc_start; sc <= sc_end; ++sc) { |
| 56 | auto sidx = n * C * H * W + sc * H * W + hw; |
| 57 | float res = sptr[didx] * -beta * std::pow(tmp, -beta - 1.0f) * 2.0f * |
| 58 | sptr[sidx] * alpha; |
| 59 | if (sc == dc) |
| 60 | res += std::pow(tmp, -beta); |
| 61 | gptr[sidx] += T(res * hptr[didx]); |
| 62 | } |
| 63 | } |
| 64 | } |
| 65 | } |
| 66 | |
| 67 | } // anonymous namespace |
| 68 |
no test coverage detected