Batchnorm channels last

Trinayan_Baruah: <p>Hi,</p>
<p>I was just implementing a simple 2d batchnorm and wanted to use channels last format. When I use the code as pasted below, my GPU profiler NSight shows the forward kernels using the channels last format as indicated by their names. But the backward uses the basic batchnorm_backward_reduce kernel instead of what is expected to be used i.e batch_norm_backward_reduce_channels_last_kernel. Is there something wrong with my initialization or is this expected behavior?</p>
<p>import torch<br>
import torch.nn as nn<br>
import torch.cuda.profiler as profiler<br>
import pyprof<br>
import argparse<br>
pyprof.init()<br>
import torch.cuda.profiler as profiler</p>
<p>parser = argparse.ArgumentParser(description=‘Profile a batch normalization layer.’)<br>
parser.add_argument(‘n’, type=int)<br>
parser.add_argument(‘c’, type=int)<br>
parser.add_argument(‘h’, type=int)<br>
parser.add_argument(‘w’, type=int)</p>
<p>args = parser.parse_args()</p>
<p>input_tensor = torch.rand(args.n, args.c, args.h, args.w).half().cuda()-0.5<br>
input_tensor.requires_grad=True<br>
input_tensor = input_tensor.to(memory_format=torch.channels_last)</p>
<p>batchnorm = nn.BatchNorm2d(args.c).half().cuda()<br>
loss_fxn = nn.MSELoss()<br>
target_tensor = torch.rand(input_tensor.shape).half().cuda()<br>
target_tensor = target_tensor.to(memory_format=torch.channels_last)</p>
<p>print(“Profiling input tensor:”,input_tensor.shape)<br>
print(“On layernorm layer:”, batchnorm)</p>
<p>for i in range(25):<br>
output = batchnorm(input_tensor).to(memory_format=torch.channels_last)<br>
loss = loss_fxn(output,target_tensor)<br>
loss.backward()</p>

ptrblck: <p>I guess <code>at::native::batch_norm_backward_reduce_kernel</code> is used as you are reducing the loss, so the gradients are not passed as a 4D channels-last tensor.<br>
If you use <code>output.backward(torch.ones_like(output))</code> you would get:</p>
<pre><code class="lang-python">at::native::batch_norm_backward_reduce_channels_last_kernel
at::native::batch_norm_backward_elemt_channels_last_kernel
</code></pre>

Trinayan_Baruah: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> : Sorry. Its not clear why transforming the output using ones_like helps in this case. Won’t that convert the tensor into 1’s instead of retaining the original output?.</p>

Trinayan_Baruah: <p>Sorry. My bad. I guess what you were saying is we need to explicitly pass the gradient to backward which can be initialized to ones. Thanks.</p>

Trinayan_Baruah: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> : At the same time, we normally call backward on the loss instead of the output correct?. Does it mean that in a real network for example, we will see the backward pass of batchnorm not applying the channels last format?</p>

ptrblck: <aside class="quote no-group" data-username="Trinayan_Baruah" data-post="5" data-topic="135671">
<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/trinayan_baruah/48/28027_2.png" class="avatar"> Trinayan_Baruah:</div>
<blockquote>
<p>At the same time, we normally call backward on the loss instead of the output correct?.</p>
</blockquote>
</aside>
<p>Yes, that’s right.</p>
<aside class="quote no-group" data-username="Trinayan_Baruah" data-post="5" data-topic="135671">
<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/trinayan_baruah/48/28027_2.png" class="avatar"> Trinayan_Baruah:</div>
<blockquote>
<p>Does it mean that in a real network for example, we will see the backward pass of batchnorm not applying the channels last format?</p>
</blockquote>
</aside>
<p>No, if you add another layer (e.g. <code>nn.Conv2d</code>) you would see the channels-last kernels again. The issue with your code would be the reduction, which would drop the memory layout, if I’m not mistaken.</p>