I have reimplemented BatchNorm1D based on the implementation provided by @ptrblck (greatly appreciated!), here: https://github.com/ptrblck/pytorch_misc/blob/master/batch_norm_manual.py 37
In order to verify identical behaviour with the nn.BatchNorm equivalent, I initiate 2 models (as well as 2 optimizers), one using MyBatchNorm and one using nn.BatchNorm. Although I can verify that in the beginning, “compareBN” returns identical results (Max diff:  tensor(8.5831e-06, device='cuda:0', grad_fn=<MaxBackward1>)), as the training progresses, very tiny differences begin to emerge.
A few iterations in:
image1659×558 95.6 KB
More iteratations later:
image1657×473 71.2 KB
By the end of the 25 epochs, both the models and the results differ significantly. Note that I have set all random seeds to 0 (including deterministic behaviour for cudnn and cuda).
This reminded me of fp16 vs fp32 training. Can it be the case that the underlying implementation of nn.BatchNorm uses different precision? Or is it something else? The MyBatchNorm method is also slightly slower (+2 seconds on a 10 second epoch), but I’d assume that this is due to C++ being used for nn.BatchNorm. I used a 2 layer FCNN with 2 BatchNorm layers on MNIST. Interestingly, the 2nd BatchNorm layer is slower in showing disrepancies.