Hi,
I’m wanting to modify the PyTorch C/C++ source code for Batch (and Group, Layer, etc.) Norm layers for part of my research, which could hopefully result in a contribution to PyTorch if successful and the work is substantial.
I’m wondering what files I should look at for modifying? The hope is I can do something like nn.BatchNorm2d(num_features, my_extra_parameter=True) and it would pass my extra parameter down to ultimately where the normalization and (optional) affine transformation is computed in the code. If my_extra_parameter == False then it should behave like the normal batch norm already implemented in PyTorch, if True it should behave differently, following my modifications. I’m wanting to make this work for CPU and GPU (mainly care about GPU since this would be used in training). What I’m trying to do involves already built-in PyTorch functions like abs, sign, etc. which I would think should make making the modifications to the forward and backward functions of batch norm relatively easy. I just don’t know where to look.
Thank you!