Hi,
I am using your implementation for my experiments. On an RTX 3090, using your implementation drops the training speed by about 30%. My own implementation, one that does not derive from nn.BatchNorm2d but rather simply from nn.Module, also had a similar drop in speed. I also tried other manual implementations 6 with the same 30% drop in training performance. I then wrote a C++ extension using PyTorch’s C++ libs and ATen. It too had the same 30% drop in training speed. I have not tried writing a CUDA kernel for it though.
I am really fascinated by the default implementation. I have a couple of queries regarding this, if you can explain. What makes the default implementation so fast? How can I have a manual implementation that is as fast (or at least within 5-10% of the speed of the default implementation)?
The reason why I am asking this is actually very important. Here in the research community, we almost always write custom modules, like weight standardization, various other ways of custom normalization schemes, or some other ops. In almost all cases, we see a significant drop in speed with our custom modules. Training a vanilla ResNet is so fast, and yet making just a small change to architecture (by introducing a custom layer), drops the training speed significantly. Not only it takes longer for our models to train (thus slowing down research), but it also keeps the GPUs quite occupied (with high utilizations and power draw). Forgive me for not knowing how the default modules of PyTorch work or how they are so fast.
Thanks.