Training two deep learning models simultaneously

username: Mriganka_Nath
cooked: <p>So I have to train two models simultaneously, where the input of the 2nd model is the output of the 1st model.<br>
And the output of my 2nd model is the grad of 1st model’s output wrt its input</p>
<p>Previously I was training these 2 models independently</p>
<ol>
<li>getting the outputs from the 1st model, then calculating its grad wrt input,</li>
<li>then feeding these outputs of 1st model to my 2nd model, and checking if the produced grad from 2nd model is equal to the original grads computed at the 1st step, using MSE loss.</li>
</ol>
<p>Now I want to train all this simultaneously. And I cant think of how</p>
<p>I was calculating the grads like this</p>
<pre><code class="lang-auto">input.requires_grad = True 
opx = model(input)  
opx.backward(torch.ones_like(opx))
dx1 = input.grad.data
</code></pre>
<p>Help will be much appreciated.</p>

username: tom
cooked: <p>You could just do</p>
<pre><code class="lang-python">opt = torch.optim.Adam(list(model1.parameters()) + list(model2.parameters()), lr=3e-4)
...
    opt.zero_grad()
    out1 = model1(input)
    out2 = model2(out1)
    loss = loss(out2, target)
    loss.backward()
    opt.step()
</code></pre>
<p>(Or even stick the two models in a sequential if you want…)</p>
<p>This will cause loss.backward() to backpropagate through both model2 and model1’s computation.</p>
<p>Best regards</p>
<p>Thomas</p>

username: Mriganka_Nath
cooked: <p>Thanks, Thomas,<br>
Sorry for another question, but the problem is that the target for out2 is the grad of out1 wrt its input.<br>
so how to proceed then<br>
(I know you are thinking why am I am doing this weird thing, right now I can’t share, but if successful to implement, I would be glad to. )</p>

username: tom
cooked: <p>Ah, sorry, I got that wrong here.<br>
To not give the next bad advice: So do you have a separate loss function for model1’s ouput and it is only model2 that should be optimized to minimize the MSE between gradients and its outputs? (Vaguely rings a bell that people did something like that a years ago, but I cannot find any references now.)</p>
<p>Best regards</p>
<p>Thomas</p>

username: Mriganka_Nath
cooked: <p>Yes, I have separate loss functions.</p>
<blockquote>
<p>So do you have a separate loss function for model1’s ouput and it is only model2 that should be optimized to minimize the MSE between gradients and its outputs?</p>
</blockquote>
<p>Yes, that’s somewhat I want to do, but want to optimize both the models simultaneously.</p>

username: tom
cooked: <p>Then I would try something along the lines that you suggested.</p>
<p>Usually I advise against using <code>.backward</code> for anything except model parameters (and use <code>torch.autograd.grad</code> instead, but this seems to be the exception to the rule because we don’t want to duplicate the work of doing grad and backward.</p>
<p>So here I would do somethign like</p>
<pre><code class="lang-python">input.requires_grad_()  # I prefer this to avoid requeires_grad = True
# if you ever re-use inputs
if input.grad is not None:
    with torch.no_grad():
        input.grad.zero_()
out1 = model1(input)
loss1 = loss_fn1(out1, target)
loss1.backward()  # populates grad of model1.parameters() and input

out2 = model2(input.detach())
loss2 = mse_loss(out2, input.grad)   # no need to use .data it will not require grad unless you use create_graph=True (which you should not)

loss2.backward()  # will populate model2.parameters() gradients  
</code></pre>
<p>Then you’d want 1 or 2 optimizers…</p>
<p>Now if you indeed want the gradient w.r.t. <code>out1.sum()</code> rather than loss1, you would want to use <code>gr_in, = torch.autograd.grad(out1.sum(), input, retain_graph=True)</code> before loss1.backward.</p>
<p>Is this approximately what you wanted?</p>
<p>Best regards</p>
<p>Thomas</p>

username: Mriganka_Nath
cooked: <p>Thanks, I think this is what I was searching for<br>
Still, I have some silly doubts (I am still learning)</p>
<ol>
<li>
<p>For the optimizer part, I am using 2 optimizers, so should I use optimizer.step() and optimizer.zero_grad() after calling loss.backward() for each model or do it in the end.</p>
</li>
<li>
<p>When we do this <code>input.requires_grad_()</code> what’s the difference between input.requires_grad = True</p>
</li>
<li>
<p>like suppose in the 2nd epoch, if the input.grad is still on, so if we find grad again wont it be grad of grads .</p>
</li>
<li>
<p>Can you explain more why we are doing this</p>
</li>
</ol>
<pre><code class="lang-auto">if input.grad is not None:
    with torch.no_grad():
        input.grad.zero_()
</code></pre>
<p>Thanks, I am very grateful.</p>

username: tom
cooked: <aside class="quote no-group" data-username="Mriganka_Nath" data-post="7" data-topic="126854">
<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/mriganka_nath/48/29859_2.png" class="avatar"> Mriganka_Nath:</div>
<blockquote>
<ol>
<li>For the optimizer part, I am using 2 optimizers, so should I use optimizer.step() and optimizer.zero_grad() after calling loss.backward() for each model or do it in the end.</li>
</ol>
</blockquote>
</aside>
<p>I don’t think it matters much as long as for each model you keep the<br>
<code>zero_grads</code> → <code>backward</code> → <code>step</code> order.</p>
<aside class="quote no-group">
<blockquote>
<ol start="2">
<li>When we do this <code>input.requires_grad_()</code> what’s the difference between input.requires_grad = True</li>
</ol>
</blockquote>
</aside>
<p>Mainly the looks and that it throws an error when you misspell <code>requires_grad_</code> rather than assigning some useless field.</p>
<aside class="quote no-group">
<blockquote>
<ol start="3">
<li>like suppose in the 2nd epoch, if the input.grad is still on, so if we find grad again wont it be grad of grads .</li>
</ol>
</blockquote>
</aside>
<p>The typical thing is that you get (batches of) your inputs from a dataloader. In this case the tensors with the inputs are re-created in every epoch, so no need to worry about it.</p>
<aside class="quote no-group">
<blockquote>
<ol start="4">
<li>Can you explain more why we are doing this</li>
</ol>
<pre data-code-wrap="python"><code class="lang-python">if input.grad is not None:
    with torch.no_grad():
        input.grad.zero_()
</code></pre>
</blockquote>
</aside>
<p>I don’t think you need this in most cases (you could add a <code>print</code> statement to see if this ever is the case).<br>
If you were to re-use inputs for whatever reason, this does the equivalent of <code>optim.zero_grad()</code> but for the <code>input</code>. (Because it isn’t in any optimizer.)</p>
<p>Best regards</p>
<p>Thomas</p>