Generating one model's parameters with another model

jensw: <p>I’m trying to generate one model’s parameters (<code>ActualModel</code>) with another model (<code>ParameterModel</code>), but running into problems with autograd when I backpropagate multiple times.</p>
<p>Here’s an example <code>ActualModel</code>, but this is supposed to be generic:</p>
<pre><code class="lang-auto">class ActualModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 1, 3)
    def forward(self, x):
        return self.conv(x)
</code></pre>
<p>The <code>ParameterModel</code> wraps the <code>ActualModel</code>, freezes its parameters and in the <code>forward()</code> method, first computes the parameters, sets the values inside <code>ActualModel</code> and then computes <code>ActualModel.forward()</code>:</p>
<pre><code class="lang-auto">class ParameterModel(torch.nn.Module):
    def __init__(self, actual_model):
        super().__init__()
        # Save actual_model in a list so its parameters aren't auto-registered
        self.actual_model = [actual_model]
        # Compute total number of parameters of the actual_model
        num_parameters = np.sum([p.numel() for p in actual_model.parameters()])
        # Freeze parameters of actual_model
        for p in actual_model.parameters():
            p.requires_grad_(False)
        # Simple model to generate num_parameters parameters
        self.linear = torch.nn.Linear(1, num_parameters)
    def forward(self, param_input, actual_input):
        # Compute the parameters for the actual model
        param_output = self.linear(param_input)
        # Set the parameters in the actual model
        idx = 0
        for p in self.actual_model[0].parameters():
            # p.set_(x) derivative is not implemented, use p *= 0, p += x instead
            p.zero_().add_(param_output[idx:idx+p.numel()].view(p.shape))
            idx += p.numel()
        # Compute the output of the actual model
        return self.actual_model[0](actual_input)
</code></pre>
<p>However, when I use this in a simplified training loop, on the second epoch I get an error about backwarding through the graph a second time, even though it should be a new graph in the second epoch.</p>
<p>Training loop:</p>
<pre><code class="lang-auto">actual_model = ActualModel()
wrapped_model = ParameterModel(actual_model)
optim = torch.optim.Adam(wrapped_model.parameters(), lr=0.07)

for i in range(2):
    print(f'Epoch {i}')
    wrapped_model.zero_grad()
    output = wrapped_model(torch.ones(1), torch.ones(1, 1, 3, 3))
    loss = output - torch.ones(1, 1, 3, 3)
    loss.sum().backward()
    optim.step()
</code></pre>
<p>Output with error message:</p>
<pre><code class="lang-auto">Epoch 0
Epoch 1
RuntimeError: Trying to backward through the graph a second time, 
but the saved intermediate results have already been freed. Specify 
retain_graph=True when calling backward the first time.
</code></pre>
<p>Any help on what I’m doing wrong here would be much appreciated.</p>

Rahul_Chand: <p><a class="mention" href="/u/jensw">@jensw</a></p>
<p>This is a toy code (which is exactly similar to your example, but much simpler for the sake of sanity!).</p>
<pre><code class="lang-auto">class TmpModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.x = [nn.Parameter(torch.Tensor(1))]
        self.x[0].requires_grad = False
        self.y = nn.Parameter(torch.Tensor(1))
        self.z = nn.Parameter(torch.Tensor(1))
        
    def forward(self, i1):

        o1 = self.y * i1
        self.x[0].zero_().add_(o1)
        return self.x[0] * self.z
        
for i in range(2):
    i1 = torch.ones(1)
    out = model_new(i1)
    out.backward()
</code></pre>
<p>If you traverse the autograd graph (by repeatedly using <code>out.grad_fn.next_functions</code>). You will find the below autograd graph. First image is the computation graph for <code>iteration 0</code> &amp; second image is the computation graph for <code>iteration 1</code>.</p>
<p><em><strong>Computation graph 1</strong></em><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/b/0b522cfa4a2512e871258954215ef4ab10194282.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/0b522cfa4a2512e871258954215ef4ab10194282" title="graph1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b522cfa4a2512e871258954215ef4ab10194282_2_690x309.jpeg" alt="graph1" data-base62-sha1="1C9jom2HjGMOZ985js6xX3WfUUa" width="690" height="309" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0b522cfa4a2512e871258954215ef4ab10194282_2_690x309.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/0/b/0b522cfa4a2512e871258954215ef4ab10194282.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/b/0b522cfa4a2512e871258954215ef4ab10194282.jpeg 2x" data-dominant-color="FCFCFD"><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">graph1</span><span class="informations">808×363 19.3 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>
<p><em><strong>Computation graph 2</strong></em><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/0/70ffcb3f7d72d7906fb3ac7efa7f12f3523cf813.png" data-download-href="https://discuss.pytorch.org/uploads/default/70ffcb3f7d72d7906fb3ac7efa7f12f3523cf813" title="graph2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/0/70ffcb3f7d72d7906fb3ac7efa7f12f3523cf813_2_690x388.png" alt="graph2" data-base62-sha1="g7DvMIynVcw5jTiymmNNkvGzbrR" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/0/70ffcb3f7d72d7906fb3ac7efa7f12f3523cf813_2_690x388.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/0/70ffcb3f7d72d7906fb3ac7efa7f12f3523cf813_2_1035x582.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/0/70ffcb3f7d72d7906fb3ac7efa7f12f3523cf813.png 2x" data-dominant-color="FDFDFA"><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">graph2</span><span class="informations">1280×720 16.2 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>
<p>The <strong>dashed “red” box</strong> has the old nodes of <code>i1</code> and <code>self.y</code>. Why? This is due to how in-place operation combines with <code>autograd</code> (see this post : <a href="https://discuss.pytorch.org/t/how-in-place-operation-affect-autograd/31173/2">link</a> ). Since the <code>computation graph 2</code> has old values as one of the nodes, you get the <code>RunTime</code> error.</p>
<p>Just for reference, the things I used to debug this was using <code>.grad_fn.next_functions</code> &amp; running the code within a <code>with torch.autograd.set_detect_anomaly(True)</code> block. You can ping me if you need any more details regarding debugging. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

jensw: <p><a class="mention" href="/u/rahul_chand">@Rahul_Chand</a> Thank you very much for the detailed analysis and for explaining how you were able to debug the issue! I understand now where the problem comes from, the <code>zero_</code> inplace operation does not make autograd “forget” the previous history even though it will no longer be relevant for computing the gradient value.</p>
<p>In your simplified example, I could fix this problem by replacing the in-place operation with the following: <code>self.x[0] = o1</code></p>
<p>However, I have not been able to to transfer this to my version where I have a wrapped model with multiple <code>Parameter</code>s inside where I don’t know the exact structure (wanting to keep this generic), so can only access it through either <code>actual_model.parameters()</code> or <code>actual_model.state_dict()</code>.</p>
<p>Do you have any idea how I could achieve this?</p>