Decomposing a conv into two smaller convs doesn't have the same grad

tayaranian: <p>I am trying to break down a conv2d into two smaller convs along the <em>out_channels</em> dimension.<br>
The output of the smaller convs are concatenated along the <em>out_channels</em> dimension to create the original output.<br>
I copy the weights from the original conv and try to compare the outputs and gradients. I observe that while the outputs are the same, the gradients have a slight error (around 2e-7). Although this error seems to be small, during a long training (e.g. Resnet18 on CIFAR-10) it causes the network to diverge.</p>
<p>Here’s the sample code for you to try. The printed diff between grads is not zero.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

device = 'cuda'

c_out = 513
c_temp = 1
conv = nn.Conv2d(2, c_out, kernel_size=2, bias=False).to(device)
conv1 = nn.Conv2d(2, c_temp, kernel_size=2, bias=False).to(device)
conv2 = nn.Conv2d(2, c_out - c_temp, kernel_size=2, bias=False).to(device)

with torch.no_grad():
    conv1.weight[:] = conv.weight[:c_temp]
    conv2.weight[:] = conv.weight[c_temp:]

inp = torch.rand(1,2,3,3).to(device)
out  = conv(inp)
out1 = conv1(inp)
out2 = conv2(inp)

cat = torch.cat((out1, out2), dim=1)

grad = torch.ones(1,c_out,2,2).to(device)
cat.backward(grad)
out.backward(grad)
print(conv.weight.grad[:c_temp] - conv1.weight.grad)
</code></pre>

KFrank: <p>Hi Mohammedreza!</p>
<p>I haven’t looked at your code, but I have a couple of comments.</p>
<aside class="quote no-group" data-username="tayaranian" data-post="1" data-topic="162981" 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/tayaranian/48/53786_2.png" class="avatar"> tayaranian:</div>
<blockquote>
<p>the gradients have a slight error (around 2e-7).</p>
</blockquote>
</aside>
<p>This is the size of round-off error for single-precision numbers (<code>torch.float32</code>),<br>
so this is to be expected.</p>
<aside class="quote no-group">
<blockquote>
<p>Although this error seems to be small, during a long training (e.g. Resnet18 on CIFAR-10) it causes the network to diverge.</p>
</blockquote>
</aside>
<p>By “diverge” do you mean that when training your original and composed<br>
versions drift away from one another, with the difference between them<br>
growing significantly larger, but each version training about as well as the<br>
other (even though they differ)?</p>
<p>Or do you mean that your composed version fails to train stably, with weights<br>
and / or gradients diverging off to <code>inf</code>?</p>
<p>If the former, it is to be expected, as the floating-point round-off error will<br>
continue to accumulate as you train, so the two versions will differ from one<br>
another.  It’s not that one is right and the other is wrong – they just differ due<br>
to accumulating round-off error and are arguably equally valid.</p>
<p>If the latter, you either have a bug somewhere, or your training is inherently<br>
unstable, so that one version trains well, but an equivalent version that only<br>
differs due to round-off error ends up running off to infinity.</p>
<p>You can check this by rerunning your test in double precision (<code>torch.float64</code>).<br>
If due to round-off error, your initial deviation should drop to something on the<br>
order of <code>1.e-15</code>.</p>
<p>To check for stability issues, rerun your test multiple times with different<br>
random initializations.  If your training is on the edge of instability, sometimes<br>
both versions will train successfully, sometimes they will both run off to infinity,<br>
and sometimes one will succeed while the other fails.  If, contrariwise, your<br>
original version always trains successfully, while your composed version<br>
always fails, you most likely have a bug in your decomposition.</p>
<p>Best.</p>
<p>K. Frank</p>