Hi all,
I try to implement a custom version of batch norm from scratch (to use a complex-valued network, I already have the other components, just need the batch norm).
For now, I want to write it in pure python using Module + Function:
https://pytorch.org/docs/stable/notes/extending.html 26
To help me, I read the way batch norm is coded in:
https://pytorch.org/docs/stable/_modules/torch/nn/modules/batchnorm.html 161
and
https://pytorch.org/docs/stable/_modules/torch/nn/functional.html 44
The module is defined in torch.nn.modules.batchnorm, where running_mean and running_var are created as buffers and then passed to the forward function that called nn.functional.batch_norm that takes running_mean as argument, itself calling the cpp batchnorm function.
But at no point running_mean is updated in the python code I read, how does the Module keeps track of it if it is only evaluated in the cpp code? What would be the best way to do it in pure python?
Thanks for your help,
SMP