fix(aten::batch_norm): A new batch norm implementation that hopefully doesnt have the same performace cost

[narendasan]

Signed-off-by: Naren Dasan naren@narendasan.com
Signed-off-by: Naren Dasan narens@nvidia.com

Description

Addresses performance issues seen with large input sizes and the conv based batch norm implementation. This new implementation just uses scale layers instead of conv.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation and have regenerated the documentation (make html in docsrc)
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes

[narendasan]

Fix was confirmed

[JIT]: batch_size: 1
    Average latency: 61.2975 ms
    Average FPS: 16.3139 fps
    Latency Standard Deviation: 0.805419
    FPS Standard Deviation: 0.24558
(excluding initial warmup runs)
[JIT/TRT]: batch_size: 1
    Average latency: 31.1957 ms
    Average FPS: 32.0557 fps
    Latency Standard Deviation: 0.107248
    FPS Standard Deviation: 0.110173