(x, w, b, affine, eps, group, fmt)
| 215 | |
| 216 | |
| 217 | def group_norm(x, w, b, affine, eps, group, fmt): |
| 218 | assert x.ndim == 4, f"group/instance norm requires 4D input, get {x.shape}" |
| 219 | N = x.shape[0] |
| 220 | |
| 221 | if fmt == "NCHW": |
| 222 | grouped_shape = (N, group, -1) |
| 223 | affine_shape = (1, -1, 1, 1) |
| 224 | norm_axis = 2 |
| 225 | elif fmt == "NHWC": |
| 226 | # grouped_shape = (N, -1, group) |
| 227 | # affine_shape = (1, 1, 1, -1) |
| 228 | # norm_axis = 1 |
| 229 | assert False, "NHWC format is not supported" |
| 230 | else: |
| 231 | assert False, f"invalid format: {fmt}" |
| 232 | |
| 233 | # reshape (N, C, H, W) to (N, Group, -1) and do normalize for each group |
| 234 | grouped_x = x.reshape(grouped_shape) |
| 235 | y, x_mean, rstd = _normalize(grouped_x, norm_axis, eps) |
| 236 | y = y.reshape(x.shape) # (N, Group, -1) -> (N, C, H, W) |
| 237 | |
| 238 | if affine: |
| 239 | w = w.reshape(affine_shape) # (C,) -> (1, C, 1, 1) |
| 240 | b = b.reshape(affine_shape) # (C,) -> (1, C, 1, 1) |
| 241 | y = y * w + b |
| 242 | |
| 243 | return y, x_mean, rstd |
| 244 | |
| 245 | |
| 246 | def group_norm_grad(dy, x, w, x_mean, rstd, affine, group, fmt): |
no test coverage detected