Hi again,
I have a few remarks about your code:
Line 63, I guess it should be n = input.numel() / input.size(1)
Then for the calculation of the variance, var = input.var([0, 2, 3], unbiased=False) gives me:
TypeError: var(): argument 'dim' (position 1) must be int, not list
I have pytorch 1.1.0a0+ac206a9, maybe it is version related, however, using
var = input.transpose(1,0).contiguous().view(input.size()[1],-1).var(-1, unbiased=False)
did the trick for me.
By running your test, the max difference between the parameters of the stock batchnorm and the custom one is about 1e-6, I’d say good enough for me. I did not quantify the time difference yet, but I ran a network with it, and it is significantly slower, but I do not mind right now.
I implemented the complex version following your example and this paper 3, it should work, but I am a bit worry about the speed now at it is heavier than the real version. The ball is in my court now.
Thanks again,
SMP