Optimizing a model in two steps with full and partial parameter updates

SarthakYadav: <p>Hi,</p>
<p>I’m trying to train a model that’s composed of two submodules, B and D, on input I.<br>
Let X_b be the intermediate feature vectors after applying B to I. (X_b = B(I))<br>
Now, D is applied on X_b to get the prediction, on which loss L_s is computed.</p>
<p>What I’m trying to do is as follows:</p>
<ol>
<li>Step 1: do the above-mentioned forward pass through the entire model, get the loss and propagate gradients backward throughout both B and D.</li>
<li>Step 2: Apply a augmentation/permutation f() on X_b to get X_b’, do a forward call on the submodule D with X_b’ and compute the same loss, but only update either B or D.</li>
</ol>
<p>I’m attaching an outline sketch of the process. Solid lines of the same colour represent a forward, and dashed lines represent backward gradient updates.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/e/be450764800133168173632da54740da3b446676.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/be450764800133168173632da54740da3b446676" title="Outline_240126_021357"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be450764800133168173632da54740da3b446676_2_690x488.jpeg" alt="Outline_240126_021357" data-base62-sha1="r9cC0Epdl1lX47ZwIk01yQmtGZg" width="690" height="488" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be450764800133168173632da54740da3b446676_2_690x488.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be450764800133168173632da54740da3b446676_2_1035x732.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/e/be450764800133168173632da54740da3b446676_2_1380x976.jpeg 2x" data-dominant-color="FBFAFA"><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">Outline_240126_021357</span><span class="informations">2262×1600 111 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>What is the best way to achieve this? I’ve tried quite a few things, including trying to set requires_grad=False for submodule D after the Step 1, but all permutations of things I have tried end up not updating the weights of D at all.</p>

KFrank: <p>Hi Sarthak!</p>
<aside class="quote no-group" data-username="SarthakYadav" data-post="1" data-topic="196088" 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/sarthakyadav/48/67251_2.png" class="avatar"> SarthakYadav:</div>
<blockquote>
<ol start="2">
<li>Step 2: Apply a augmentation/permutation f() on X_b to get X_b’, do a forward call on the submodule D with X_b’ and compute the same loss, but only update either B or D.</li>
</ol>
</blockquote>
</aside>
<p>You don’t use the word “optimize” in your description, but I assume that you<br>
want to backpropagate the loss <code>L_s</code> in order to populate the gradients of <code>B</code><br>
and <code>D</code> and then take an optimization step.</p>
<p>You make the logic a little more complicated by saying in Step 2 that you<br>
want to only update either B or D.</p>
<p>Let’s start with the case where you update D in Step 2:</p>
<p>Perform your Step 1 forward pass, computing <code>L_s</code> and keeping a reference<br>
to <code>X_b</code>.  Make a copy of <code>X_b</code>, <code>X_c = X_b.detach().clone()</code>, pass <code>f (X_c)</code><br>
through D, and compute a second value for the loss <code>L_s</code>.  Sum the two<br>
losses together and backpropagate.</p>
<p>The will populate B with gradients due to the Step 1 forward pass and will<br>
populate D with the sum of its Step 1 and Step 2 forward-pass gradients.<br>
You can now optimize with <code>opt_B.step()</code> and <code>opt_D.step()</code>.</p>
<p>For the case where you update B in Step 2:</p>
<p>Perform the forward pass just through B.  Keep a reference to <code>X_b</code>  Pass<br>
<code>f (X_b</code>) through D (no <code>.detach()</code> nor <code>.clone()</code>) and compute <code>L_s</code>.  Call<br>
<code>L_s (retain_graph = True)</code>.  This populates B with its Step 2 gradients.<br>
(It also populates D with undesired gradients.)</p>
<p>Now call <code>opt_D.zero_grad()</code>.  This zeros out the Step 2 gradients for D<br>
that you said you didn’t want.  Complete the Step 1 forward pass by passing<br>
<code>X_b</code> (to which you kept a reference) through D, compute <code>L_s</code>, and call<br>
<code>L_s.backward()</code>.  This populates D’s Step 1 gradients and accumulates<br>
B’s Step 1 gradients onto its Step 2 gradients.</p>
<p>Now optimize with <code>opt_B.step()</code> and <code>opt_D.step()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>