nn.BatchNorm is slow

username: dailaila
cooked: <p>Referring to my previous <a href="https://discuss.pytorch.org/t/nn-conv-is-slower-than-before/212053">question</a> about a custom convolution layer, I figured out that the slowness may not be due to the convolution operation, but rather to the batch normalization applied after that.</p>
<p>As you can see here (and also in the old question), takes 76% of the computation time:</p>
<pre><code class="lang-auto">with torch.autograd.profiler.profile(use_cuda=True, record_shapes=True) as prof:
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1517/1517 [09:27&lt;00:00,  2.67it/s] 
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                       aten::batch_norm         0.13%     559.534ms         3.79%       15.902s     698.815us     264.701ms         0.05%      429.567s      18.878ms         22755
                           aten::_batch_norm_impl_index         0.12%     516.742ms         3.66%       15.342s     674.225us     382.006ms         0.07%      429.302s      18.866ms         22755
                                 aten::cudnn_batch_norm         2.37%        9.968s         3.48%       14.618s     741.245us      425.741s        76.61%      428.051s      21.705ms         19721
                                             aten::topk         0.16%     682.746ms         0.16%     682.746ms     112.516us       20.922s         3.76%       20.922s       3.448ms          6068
                                      aten::convolution         0.27%        1.115s         3.41%       14.329s     674.702us      59.594ms         0.01%       18.548s     873.326us         21238
                                     aten::_convolution         2.57%       10.804s         3.15%       13.215s     622.223us     253.331ms         0.05%       18.488s     870.520us         21238
                                       aten::leaky_relu         0.12%     491.424ms         0.12%     491.424ms      21.596us       17.673s         3.18%       17.673s     776.669us         22755
                                              aten::max         0.08%     355.966ms         0.09%     368.209ms      40.454us       17.145s         3.09%       17.165s       1.886ms          9102
                                aten::cudnn_convolution         0.27%        1.141s         0.27%        1.141s      53.743us       17.073s         3.07%       17.073s     803.882us         21238
                                           aten::conv2d         0.03%     116.006ms         0.18%     748.824ms      70.517us      28.624ms         0.01%       15.054s       1.418ms         10619
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 419.744s
Self CUDA time total: 555.749s
</code></pre>
<p>So now I was wondering what causes this. My batch size isn’t that big, I only have <code>8</code> as batch size and the embedding size is nearly always around <code>64</code>, while the remaining dimensions (are fixed and) never exceed <code>2048</code> and <code>40</code>. So my typical comvolution input is around <code>8x64x2048x40</code>.</p>