BatchNorm breaks down with tiny batches (common when training on high-resolution images). GroupNorm (Wu & He, 2018) normalises each sample on its own, over groups of channels.
For x of shape (N, C, H, W) and G groups (C divisible by G): reshape to (N, G, C/G, H, W), normalise over the last three axes with the population variance and eps, reshape back, then apply a per-channel scale gamma and shift beta (shape (C,)).
Implement group_norm(x, gamma, beta, groups, eps=1e-5). Two special cases are worth knowing:
groups = 1 is LayerNorm over (C, H, W).groups = C is InstanceNorm.