Pixel normalization through channels

username: patrickctrf
cooked: <p>I’m trying to reproduce <a href="https://arxiv.org/pdf/1902.08710.pdf" rel="noopener nofollow ugc">GANSynth</a> paper and they make use of Pixel Norm, a technique that computes the norm of a single pixel among all channels per sample:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/3/d3218ef0f289da587bb81423b5b710139f7afd0c.png" data-download-href="https://discuss.pytorch.org/uploads/default/d3218ef0f289da587bb81423b5b710139f7afd0c" title="Captura de tela de 2022-07-05 15-01-38"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3218ef0f289da587bb81423b5b710139f7afd0c_2_690x212.png" alt="Captura de tela de 2022-07-05 15-01-38" data-base62-sha1="u7KDYkM3iSrt9Vpzq1ylX21Pf1q" width="690" height="212" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3218ef0f289da587bb81423b5b710139f7afd0c_2_690x212.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3218ef0f289da587bb81423b5b710139f7afd0c_2_1035x318.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/3/d3218ef0f289da587bb81423b5b710139f7afd0c.png 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Captura de tela de 2022-07-05 15-01-38</span><span class="informations">1050×323 45 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<hr>
<p>How can I have this kind of normalization in PyTorch?</p>
<p>I tried torch.nn.functional.normalize(input, dim=1), but it normalizes through all batch samples at once. Pixel norm intends to normalize inside a single sample (similar to LayerNorm).</p>

username: ptrblck
cooked: <p>I think you could try to directly implement the normalization via:</p>
<pre><code class="lang-python">b, c, h, w = 2, 3, 4, 4
x = torch.randn(b, c, h, w)

out = x  / (1/x.size(1) * (x**2).sum(dim=1, keepdims=True)).sqrt()
</code></pre>
<p>If you have any reference implementation you could compare these approaches and see if I misunderstood the formula.</p>

username: patrickctrf
cooked: <p>I’ll check and then post the feedback here.</p>
<p>Update: It looks correct. Thanks!</p>

username: ptrblck
cooked: <p>Great and thanks for confirming!</p>
<p>Btw., note that each of these primitive operations would launch a separate CUDA kernel (in case you are using the GPU) so you might not see the best performance.<br>
If you are using PyTorch <code>&gt;=1.12.0</code> you could try to <code>torch.jit.script</code> it and allow <code>nvFuser</code> to code generate fast kernels for your workload.</p>
<p>Here is an example which sees a ~3.8x speedup via the generated CUDA kernels:</p>
<pre><code class="lang-python">def my_norm(x):
    out = x  / (1/x.size(1) * (x**2).sum(dim=1, keepdim=True)).sqrt()
    return out

iteration_count = 100
x = torch.randn(64, 3, 1024, 1024, device='cuda')

# Eager execution
# Perform warm-up iterations
for _ in range(3):
    output = my_norm(x)

# Synchronize the GPU before starting the timer
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(iteration_count):
    output = my_norm(x)

# Synchronize the GPU before stopping the timer
torch.cuda.synchronize()
stop = time.perf_counter()
iters_per_second = iteration_count / (stop - start)
print("Average iterations per second: {:.2f}".format(iters_per_second))


# Scripted via nvFuser
my_norm_scripted = torch.jit.script(my_norm)

# Perform warm-up iterations
for _ in range(3):
    output = my_norm_scripted(x)

# Synchronize the GPU before starting the timer
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(iteration_count):
    output = my_norm_scripted(x)

# Synchronize the GPU before stopping the timer
torch.cuda.synchronize()
stop = time.perf_counter()
iters_per_second = iteration_count / (stop - start)
print("Average iterations per second: {:.2f}".format(iters_per_second))
</code></pre>
<p>Output:</p>
<pre><code class="lang-python"># Eager
Average iterations per second: 135.68

# Scripted with generated kernels from nvFuser
Average iterations per second: 515.90
</code></pre>