| 29 | |
| 30 | |
| 31 | def _np_bn_training(x, scale, bias, rm, rv, momentum=0.1, e=1e-5): |
| 32 | channel = x.shape[1] |
| 33 | np.testing.assert_array_almost_equal(scale.shape, (1, channel, 1, 1)) |
| 34 | np.testing.assert_array_almost_equal(bias.shape, (1, channel, 1, 1)) |
| 35 | np.testing.assert_array_almost_equal(rm.shape, (1, channel, 1, 1)) |
| 36 | np.testing.assert_array_almost_equal(rv.shape, (1, channel, 1, 1)) |
| 37 | |
| 38 | batch_m = x.mean(axis=(0, 2, 3), keepdims=True) |
| 39 | batch_v = x.var(axis=(0, 2, 3), keepdims=True) |
| 40 | |
| 41 | x_norm = (x - batch_m) / np.sqrt(batch_v + e) |
| 42 | y_norm = x_norm * scale + bias |
| 43 | |
| 44 | # https://arxiv.org/pdf/1502.03167.pdf |
| 45 | s = list(x.shape) |
| 46 | s[1] = 1 |
| 47 | batch_v_unbiased = np.prod(s) * batch_v / (np.prod(s) - 1) |
| 48 | |
| 49 | rm = momentum * batch_m + (1 - momentum) * rm |
| 50 | rv = momentum * batch_v_unbiased + (1 - momentum) * rv |
| 51 | |
| 52 | # https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnBatchNormalizationForwardTraining |
| 53 | resultSaveInvVariance = 1 / np.sqrt(batch_v) |
| 54 | return y_norm, rm, rv, batch_m, resultSaveInvVariance |
| 55 | |
| 56 | |
| 57 | def _np_bn_testing(x, scale, bias, rm, rv, momentum=0.1, e=1e-5): |