Invalid gradient shape after discarding filters during training

kuonlp: <p>Hello,</p>
<p>I’m trying to remove some filters during training, however, the .backward() in the second iteration raises an error due to the size mismatch. The error is the following:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "fixingbackward.py", line 44, in &lt;module&gt;
    loss.backward() # Error here
  File "/home/user/data/envs/def/lib/python3.8/site-packages/torch/tensor.py", line 222, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/user/data/envs/def/lib/python3.8/site-packages/torch/autograd/__init__.py", line 130, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Function CudnnConvolutionBackward returned an invalid gradient at index 1 - got [1, 7, 3, 3] but expected shape compatible with [1, 8, 3, 3]
</code></pre>
<p>The following code reproduces the error:</p>
<pre><code class="lang-auto">import torch

class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 8, 3, padding=1)
        self.conv2 = torch.nn.Conv2d(8, 1, 3, padding=1)

    def forward(self, x): 
        x = self.conv1(x)
        x = self.conv2(x)
        return x

def loss_fun(x):
    return torch.sum(x)

# Initialization
model = Model().cuda()
opt = torch.optim.Adam(model.parameters(), lr=1e-4)
im = torch.rand((1, 1, 256, 256)).cuda()

# Iteration 1
output = model(im)
loss = loss_fun(output)

opt.zero_grad()
loss.backward()
opt.step()

# Remove the first channel of the filters (and the grads)
model.conv1.weight.data = model.conv1.weight.data[1:] # size = (7, 1, 3, 3)
model.conv1.bias.data = model.conv1.bias.data[1:] # size = (7)
model.conv2.weight.data = model.conv2.weight.data[:, 1:] # size = (1, 7, 3, 3)
#### Alternatively, this seems to work, but it increases memory usage until it explodes
# model.conv1.weight = torch.nn.Parameter(model.conv1.weight.data[1:]) # (7, 1, 3, 3)
# model.conv1.bias = torch.nn.Parameter(model.conv1.bias.data[1:]) # (7)
# model.conv2.weight = torch.nn.Parameter(model.conv2.weight.data[:, 1:]) # (1, 7, 3, 3)

# Iteration 2
output = model(im)
loss = loss_fun(output)

opt.zero_grad()
loss.backward() # Error here
opt.step()
</code></pre>
<p>Following <a href="https://discuss.pytorch.org/t/dynamically-expandable-networks-resizing-weights-during-training/19218">this post</a>, I also tried 1) to update weight.grad.data similarly, and 2) to even initialize the optimizer right before Iteration 2. But these approaches didn’t work.</p>
<p>After debugging for a long time I’m stuck here since I don’t know where does the information about the shape comes from. In fact, I don’t understand why this problem arises. I would expect that when I do the output = model(im) autograd realizes that the shapes have changed already and should ignore whatever it knows from Iteration 1. It seems that it stores some information from the previous iteration but I have no clue what/where can be.</p>
<p>Thanks!</p>

ptrblck: <aside class="quote no-group" data-username="kuonlp" data-post="1" data-topic="126285">
<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/k/fbc32d/48.png" class="avatar"> kuonlp:</div>
<blockquote>
<p>I would expect that when I do the output = model(im) autograd realizes that the shapes have changed already and should ignore whatever it knows from Iteration 1.</p>
</blockquote>
</aside>
<p>That’s not the case, since you are manipulating the underlying <code>.data</code> attributes, which Autograd doesn’t track, and which is thus not recommended.<br>
Try to assign the new parameters directly and, if needed, wrap it into a <code>with torch.no_grad()</code> guard.</p>

kuonlp: <p>Edit: I’ve solved the problem, so I re-wrote my initial answer and what I initially thought it was the problem. Then I write the actual problem I had (which was embarrassingly simple).</p>
<p>I initially though that with torch.nn.Parameter I would be creating new parameters all the time, and because of this I run out of memory. In fact, the following code makes it run out of memory</p>
<pre><code class="lang-auto">with torch.no_grad():
    for mod in self.layer.modules():
        if isinstance(mod, torch.nn.Conv2d) or isinstance(mod, torch.nn.Conv3d    ):
            pp = mod.weight
            del mod.weight
            mod.weight = torch.nn.Parameter(pp.data)
</code></pre>
<p>whereas this doesn’t make the code run out of memory. Note: I have tested this on my actualy script, not on the toy example I gave in the initial post.</p>
<pre><code class="lang-auto">with torch.no_grad():
    for mod in self.layer.modules():
        if isinstance(mod, torch.nn.Conv2d) or isinstance(mod, torch.nn.Conv3d    ):
            pp = mod.weight
            del mod.weight
            mod.weight = pp
</code></pre>
<p>However, I examined what’s on memory (following <a href="https://discuss.pytorch.org/t/how-to-debug-causes-of-gpu-memory-leaks/6741/3">this</a>), and I noticed that the number of objects in memory (including torch.nn.Parameters) and the total size increase at first, but it becomes constant in both cases. How is it possible that despite the number of tensors, parameters, and their size are constant, the script runs out of memory?</p>
<p><strong>Solution</strong><br>
In my code I kept track of the loss in this way:</p>
<pre><code class="lang-auto">tr_loss_iteration = loss(output, Y)
tr_loss += tr_loss_iteration # For knowing the avg loss in the current epoch.
</code></pre>
<p>Apparently, tr_loss  was accumulating something, although I’m not sure what, because, as I’ve said, the number of tensors and parameters in memory was constant. Maybe it was keeping track of all “new” parameters? Anyway, I solved the memory-leaking problem by simply:</p>
<pre><code class="lang-auto">tr_loss += tr_loss_iteration.cpu().detach().numpy()
</code></pre>
<p>Thanks!</p>