BatchNorm2d
Normalise each channel over the batch and spatial axes.
Two variance estimators, deliberately. The batch is normalised with the biased variance (divide by N) while the running estimate accumulates the unbiased one (divide by N−1). That is torch's behaviour, verified, and the asymmetry is principled: the biased figure is the right normaliser for the batch in hand, the unbiased one the right estimator of the population.
Using one for both makes evaluation drift away from training as the running estimate converges to the wrong value — which a single-step comparison cannot see, so it is pinned by a test that runs many.
The running statistics are updated under no_grad and assigned in place. They are bookkeeping about the data seen so far, not part of the function being differentiated, and letting them onto the tape would keep every past batch's graph alive.
A single-sample batch leaves the running estimate untouched rather than dividing by zero, matching torch.
num_batches_tracked exists and is deliberately not maintained. Every torch BatchNorm state_dict carries the key, so the buffer must exist or load_state_dict rejects the checkpoint — but nothing in vkML reads it, since torch uses it only for momentum=None, a mode this layer does not offer.
Keeping it accurate would cost a host round-trip per forward pass: int64 arithmetic is unimplemented on both backends, so the increment cannot happen on the device, and reading the counter back is exactly the per-step synchronisation this project spends effort avoiding.
Stated consequence: a vkML checkpoint loaded into torch reports zero batches tracked.
Construction¶
__init__¶
Forward¶
forward¶
Internals¶
__repr__¶
See also batch_norm, LayerNorm, Module