MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / _normalize

Function _normalize

imperative/python/megengine/xla/rules/normalize.py:124–147  ·  view source on GitHub ↗

normalize X along axes, return normalized, mean, rstd for example: x.shape = (N, C, H, W), axes = (1, 2, 3) return normalized(N, C, H, W), mean(N,), rstd(N)

(x, axes, eps)

Source from the content-addressed store, hash-verified

122
123
124def _normalize(x, axes, eps):
125 """
126 normalize X along axes, return normalized, mean, rstd
127 for example:
128 x.shape = (N, C, H, W), axes = (1, 2, 3)
129 return normalized(N, C, H, W), mean(N,), rstd(N)
130 """
131 axes = _normalize_reduce_axes(x.shape, axes)
132 x_mean = x.mean(axes, True)
133 x_mean_sqr = x_mean * x_mean
134 x_sqr = x * x
135 x_sqr_mean = x_sqr.mean(axes, True)
136 var = x_sqr_mean - x_mean_sqr
137 var_plus_eps = var + eps
138 std = sqrt(var_plus_eps)
139 rstd = 1.0 / std
140 delta = x - x_mean
141 normalized = delta * rstd
142
143 non_reduce_shape = [x.shape[i] for i in range(x.ndim) if i not in axes]
144 x_mean = x_mean.reshape(non_reduce_shape)
145 rstd = rstd.reshape(non_reduce_shape)
146
147 return normalized, x_mean, rstd
148
149
150def _normalize_grad(dy, x, x_mean, rstd, axes):

Callers 2

layer_normFunction · 0.85
group_normFunction · 0.85

Calls 4

_normalize_reduce_axesFunction · 0.85
sqrtFunction · 0.50
meanMethod · 0.45
reshapeMethod · 0.45

Tested by

no test coverage detected