How to implement efficient element-wise 'ax+b' function

zhi_Li: <p>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.</p>
<pre><code class="lang-auto">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                   
</code></pre>
<pre><code class="lang-auto">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
</code></pre>
<p>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.</p>
<p>I work on Pytorch1.8.0a0+79b9c03 comiled with cuda11.1 and cudnn 8005.</p>

googlebot: <p>try these:</p>
<pre><code class="lang-auto">@jit.script
def affine(x,m,a):
  return x*m+a

def affine(x,m,a):
  return (x*m).add_(a)
</code></pre>
<p>also, view() overhead in your test may be non-negligible.</p>

KFrank: <p>Hi Zhi and Alex!</p>
<aside class="quote no-group" data-username="zhi_Li" data-post="1" data-topic="105274" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zhi_li/48/14873_2.png" class="avatar"> zhi_Li:</div>
<blockquote>
<p>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,</p>
</blockquote>
</aside>
<p>The short story is that I can reproduce your results.</p>
<p>I’ve also done some further experimentation.</p>
<aside class="quote no-group">
<blockquote>
<p>So what is the best way to implement the scaling and bias function?</p>
</blockquote>
</aside>
<p>The best I could come up with was to apply Alex’s torchscript suggestion<br>
to your “w-b” version.</p>
<p>I also tried a <code>Conv2d</code> version – even though it doesn’t have<br>
<code>BatchNorm2d</code>’s normalization step, it was also slower.</p>
<p>Note, torchscript gave a nice speedup with pytorch version 1.6.0, but<br>
not with a version 1.8.0 nightly build.</p>
<p>I used your original problem – same-shaped tensors, and such – but<br>
I modified the code to perform some additional tests and to tie out the<br>
results of the various methods with one another.</p>
<p>Here is the script:</p>
<pre data-code-wrap="python"><code class="lang-python"># cxb_time.py

# exec (open ('./cxb_time.py').read())

import torch
print ('torch.__version__ =', torch.__version__)
print ('torch.cuda.get_device_name (0) =', torch.cuda.get_device_name (0))

import time

# contents of torchscript file, jts.py:

# import torch
#
# @torch.jit.script
# def jtwb (w, b, x):
#     return x * w.view (x.shape[0], x.shape[1], 1, 1) + b.view (x.shape[0], x.shape[1], 1, 1)
#
# @torch.jit.script
# def trmac (w, b, x):
#     return ((w * x.transpose (1, 3)).add_ (b)).transpose (1, 3)

from jts import jtwb, trmac

N, C, H, W = 64, 256, 64, 64
repetitions = 100

# random weight and bias
w = torch.randn (C).cuda()
b = torch.randn (C).cuda()

# initialize bn, and cv identically to w-b
bn = torch.nn.BatchNorm2d(C, eps = 0.0).cuda()
cv = torch.nn.Conv2d (C, C, 1, groups = C).cuda()
with torch.no_grad():
    _ = bn.weight.copy_ (w)
    _ = bn.bias.copy_ (b)
    _ = cv.weight.copy_ (w.reshape ([C, 1, 1, 1]))
    _ = cv.bias.copy_ (b)

w0 = w.clone()  # for trmac version
b0 = b.clone()  # for trmac version
w = torch.nn.Parameter(w.expand(N, C).contiguous())
b = torch.nn.Parameter(b.expand(N, C).contiguous())

x = torch.randn(N, C, H, W).cuda()
# normalize x so that only weight and bias (gamma and beta) of batchnorm matter
x = (x - x.mean (dim = 0)) / x.std (dim = 0, unbiased = False)

# compute y five different ways
ywb = x * w.view(N, C, 1, 1) + b.view(N, C, 1, 1)
ybn = bn (x)
ycv = cv (x)
yjt = jtwb (w, b, x)
ytr = trmac (w0, b0, x)

# check that results for y are nearly equal
print ('verify that the different methods give the same result:')
print ('(ybn - ywb).abs().max() =', (ybn - ywb).abs().max().item())
print ('(ycv - ywb).abs().max() =', (ycv - ywb).abs().max().item())
print ('(yjt - ywb).abs().max() =', (yjt - ywb).abs().max().item())
print ('(ytr - ywb).abs().max() =', (ytr - ywb).abs().max().item())

ywb = None
ybn = None
ycv = None
yjt = None
ytr = None

print ('time the different methods, repetitions =', repetitions)

torch.cuda.synchronize()
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 ('w-b time:', (time.time() - t0) / repetitions)

torch.cuda.synchronize()
t0 = time.time()
for _ in range(repetitions):
    y = bn(x)

torch.cuda.synchronize()
print ('bn time:', (time.time() - t0) / repetitions)

torch.cuda.synchronize()
t0 = time.time()
for _ in range(repetitions):
    y = cv(x)

torch.cuda.synchronize()
print ('cv time:', (time.time() - t0) / repetitions)

torch.cuda.synchronize()
t0 = time.time()
for _ in range(repetitions):
    y = jtwb (w, b, x)

torch.cuda.synchronize()
print ('jtwb time:', (time.time() - t0) / repetitions)

torch.cuda.synchronize()
t0 = time.time()
for _ in range(repetitions):
    y = trmac (w0, b0, x)

torch.cuda.synchronize()
print ('jttr time:', (time.time() - t0) / repetitions)
</code></pre>
<p>Here is the pytorch version 1.6.0 output:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; exec (open ('./cxb_time.py').read())
torch.__version__ = 1.6.0
torch.cuda.get_device_name (0) = GeForce GTX 1050 Ti
verify that the different methods give the same result:
(ybn - ywb).abs().max() = 2.384185791015625e-07
(ycv - ywb).abs().max() = 9.5367431640625e-07
(yjt - ywb).abs().max() = 9.5367431640625e-07
(ytr - ywb).abs().max() = 0.0
time the different methods, repetitions = 100
w-b time: 0.010975425243377685
bn time: 0.008232767581939698
cv time: 0.009469003677368163
jtwb time: 0.005830504894256592
jttr time: 0.014244294166564942
</code></pre>
<p>And, for completeness, here is the 1.8.0.dev20201203 output:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; exec (open ('./cxb_time.py').read())
torch.__version__ = 1.8.0.dev20201203
torch.cuda.get_device_name (0) = GeForce GTX 1050 Ti
verify that the different methods give the same result:
(ybn - ywb).abs().max() = 4.76837158203125e-07
(ycv - ywb).abs().max() = 9.5367431640625e-07
(yjt - ywb).abs().max() = 0.0
(ytr - ywb).abs().max() = 0.0
time the different methods, repetitions = 100
w-b time: 0.011010963916778565
bn time: 0.008294436931610107
cv time: 0.009465277194976807
jtwb time: 0.011733412742614746
jttr time: 0.01428614854812622
</code></pre>
<p>The only difference is that with version 1.8.0.dev20201203 using jit<br>
does not speed up your original “w-b” method.  (Also notice, that the<br>
1.6.0 jit “w-b” method differed at round-off error from the raw “w-b”<br>
method, while the 1.8.0 jit “w-b” method exactly reproduced the raw<br>
“w-b” results.  It’s as if jit didn’t do anything.)</p>
<p>Best.</p>
<p>K. Frank</p>

googlebot: <p>You need to do two warmup iterations with 1.7+ JIT (profiling &amp; codegen)</p>

KFrank: <p>Hi Alex!</p>
<aside class="quote no-group" data-username="googlebot" data-post="4" data-topic="105274" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<p>You need to do two warmup iterations with 1.7+ JIT (profiling &amp; codegen)</p>
</blockquote>
</aside>
<p>I added the following to my timing script:</p>
<pre data-code-wrap="python"><code class="lang-python"># try "warming up" jtwb
print ('call jtwb() twice to "warm up"')
yjt1 = jtwb (w, b, x)
yjt2 = jtwb (w, b, x)
yjt1 = None
yjt2 = None
</code></pre>
<p>(So there are now three calls to <code>jtwb()</code> before the actual timing loop.)</p>
<p>This changed the results for the jtwb timing only modestly:</p>
<pre><code class="lang-plaintext">jtwb time: 0.010964112281799316
(was:      0.011733412742614746)
</code></pre>
<p>And <code>yjt</code> (the result of the <code>jtwb()</code> call) still agrees exactly with <code>ywb</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

googlebot: <p>Ok, it appears fusion optimization is fragile</p>
<ol>
<li>“from jts import *” breaks it, “import jts” works</li>
<li>creating views inside jitted function breaks it</li>
</ol>
<p>You can check this with torch.jit.last_executed_optimized_graph(), after warmup this graph should contain TensorExprGroup (or FusionGroup with legacy jit) block.</p>
<p>for the moment, it can be better to use legacy fuser:</p>
<pre><code class="lang-auto">torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_profiling_executor(False)
</code></pre>

zhi_Li: <p>Hi Alex!<br>
I think the fusion optimization of <code>torch.jit</code> only works on  inputs with fixed size that it is warmed up with (Actually I found feeding tensors with another size doesn’t work.). Is my understanding right?<br>
If I want to <strong>apply the optimized function to tensors with different sizes</strong>(e.g. <span class="hashtag">#channels</span> of feature map are increasing as going deep with a CNN), is there a way to implement it elegantly?</p>

googlebot: <p>This puzzles me a bit too - new codegen creates over-specialized kernels, hardcoding tensor sizes. However, my understanding is, it will compile multiple times for different sizes, if it sees the same shape[s] twice. There is torch._C._jit_set_bailout_depth(n) option to control this, it is 20 by default.</p>
<p>Simpler alternative is to use legacy jit, it produces more generic code.</p>