For some reason, I want to implement the scaling and bias part of BatchNorm, w/o the normalizing part. I resort to broadcasting to implement this element-wise operation. But I found such implementation is slower than the original nn.BathNorm implementation, whose theoretical FLOPs are several times more than the former.
In [49]: import torch
    ...: import time
    ...:
    ...: N, C, H, W = 64, 256, 64, 64
    ...: repetitions = 100
    ...:
    ...: x = torch.randn(N, C, H, W).cuda()
    ...: w = nn.Parameter(torch.randn(C).cuda().expand(N, C).contiguous())
    ...: b = nn.Parameter(torch.randn(C).cuda().expand(N, C).contiguous())
    ...:
    ...: t0 = time.time()
    ...: for _ in range(repetitions):
    ...:     y = x * w.view(N, C, 1, 1) + b.view(N, C, 1, 1)
    ...: torch.cuda.synchronize()
    ...: print ((time.time() - t0) / repetitions)
0.0012868165969848633                   

In [41]: import torch
    ...: import time
    ...:
    ...: N, C, H, W = 64, 256, 64, 64
    ...: repetitions = 100
    ...:
    ...: x = torch.randn(N, C, H, W).cuda()
    ...: #w = torch.randn(C).cuda().expand(N, C).contiguous()
    ...: bn = torch.nn.BatchNorm2d(C).cuda()
    ...:
    ...: t0 = time.time()
    ...: for _ in range(repetitions):
    ...:     y = bn(x)
    ...: torch.cuda.synchronize()
    ...: print ((time.time() - t0) / repetitions)
0.0009896183013916016

So what is the best way to implement the scaling and bias function?When it comes to the training phase, time per batch of the broadcasting impl. is  almost twice as the original BN impl. in my model backboned by a ResNet18.
I work on Pytorch1.8.0a0+79b9c03 comiled with cuda11.1 and cudnn 8005.