albanD:

grad_weight.sum(dim=(-1, -2))


Hi,
Thanks for checking out the code and quick reply.
I get the same size of gradient error and even when I correct it the way you suggested, I still get a malfunctioning batchnorm, meaning it gives me ‘NaN’ loss when I plug this custom bn into resnet18.
So, I added the following code to compare the gradients and I see the output backward hooks gives the same results but the inputs are different. I would very much appreciate it if you could take a look at it.
class Hook():
    def __init__(self, module, backward=False):
        if backward==False:
            self.hook = module.register_forward_hook(self.hook_fn)
        else:
            self.hook = module.register_backward_hook(self.hook_fn)
    def hook_fn(self, module, input, output):
        self.input = input
        self.output = output
    def close(self):
        self.hook.remove()
        
        
my_bn = BatchNorm2d(3, affine=True) 
bn = nn.BatchNorm2d(3, affine=True)

hookB_bn = Hook(bn,backward=True) 
hookB_mybn = Hook(my_bn,backward=True)

compare_bn(my_bn, bn)  # weight and bias should be different
# Load weight and bias
my_bn.load_state_dict(bn.state_dict())
compare_bn(my_bn, bn)

x = torch.randn(10, 3, 100, 100) * scale + bias
o2 = bn(x)
o1 = my_bn(x)

o2.backward(torch.ones_like(o1), retain_graph=True)
o1.backward(torch.ones_like(o2), retain_graph=True)

print('***'*3+'  nn.bn Backward Hooks Inputs & Outputs shapes  '+'***'*3)
             
print('input:', [i.shape for i in hookB_bn.input if hasattr(i, 'shape')])
print('output:' ,[i.shape for i in hookB_bn.output if hasattr(i, 'shape')])         
print('---'*17)

print('***'*3+'  mybn Backward Hooks Inputs & Outputs shapes  '+'***'*3)
             
print('input:',[i.shape for i in hookB_mybn.input if hasattr(i, 'shape')])
print('output:' ,[i.shape for i in hookB_mybn.output if hasattr(i, 'shape')])         
print('---'*17)

print('***'*3+'  outputs are the same  '+'***'*3)
print(hookB_mybn.output[0]==hookB_bn.output[0])

print('***'*3+'  inputs differ  '+'***'*3)
print(hookB_mybn.input[1]==hookB_bn.input[1])