Skip the calculation of gradients for some layer after forward pass

Lipril: <p>Hello,</p>
<p>Let’s say I have 6 (transformer) layers (i.e., L0, L1, L2, L3, L4, L5), all the transformer layer now have require_grads = True, Now I do a forward pass to all the layers. <strong>Based on the representations</strong>, I might don’t want to calculate the gradients for some layers (to save some time) let say (L0, L1 and L2) i.e., I want to stop the gradient calculation at L2.</p>
<p>To perform this, initially I have all the layers have require_grads = True, after forward pass, let say I don’t want to calculate the gradients for L0, L1 &amp; L2, I put the L0, L1 &amp; L2 require_grads = False. Even after doing this the program takes the same time to perform the task (I have check the gradient they are “None”, they not calculated). Can anyone help me out with this. What I am doing wrong here. Or is there any other way to do this?</p>
<p>Note: I put the word embedding layer require_grads = False.</p>
<p>Code:</p>
<p>for param in student_model.distilbert.embeddings.parameters():</p>
<p>param.requires_grad = False</p>
<p>for param in student_model.distilbert.transfomer.layer[:curr_flag].parameters():</p>
<p>param.requires_grad = False</p>
<p>Thank you in advance.</p>

KFrank: <p>Hi Lipril!</p>
<aside class="quote no-group" data-username="Lipril" data-post="1" data-topic="222185">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lipril/48/77915_2.png" class="avatar"> Lipril:</div>
<blockquote>
<p>Let’s say I have 6 (transformer) layers (i.e., L0, L1, L2, L3, L4, L5), all the transformer layer now have require_grads = True, Now I do a forward pass to all the layers. <strong>Based on the representations</strong>, I might don’t want to calculate the gradients for some layers (to save some time) let say (L0, L1 and L2) i.e., I want to stop the gradient calculation at L2.</p>
</blockquote>
</aside>
<p>If I understand your use case correctly, you want to perform a forward pass through all</p>
<p>layers, <code>L0</code>, <code>L1</code>, …, <code>L5</code>, and then based of the result of the full forward pass (that is, the</p>
<p>output of <code>L5</code>), you decide whether or not to compute gradients for <code>L0</code>, <code>L1</code>, and <code>L2</code>.</p>
<p>Since your decision depends on the full forward pass, you have to incur the cost of the</p>
<p>computing <code>L0</code>, …, <code>L2</code>, but you don’t necessarily have to incur the cost of tracking those</p>
<p>gradients.  If you <em>usually</em> will be computing gradients for <code>L0</code>, …, <code>L2</code>, you should do something</p>
<p>like:</p>
<pre><code class="lang-auto">L2_out = L2 (L1 (L0 (input)))
L2_out_detach = L2_out.detach().requires_grad_ (True)
output = L5 (L4 (L3 (L2_out_detach)))
loss = loss_fn (output)
compute_grads = should_i_compute_grads_for_l012 (output)
loss.backward()    # populates grads for L3, L4, and L5, and L2_out_detach
if  compute_grads:
    torch.autograd.backward (L2_out, grad_tensors = L2_out_detach.grad)   # populates grads for L0, L1, and L2

</code></pre>
<p>This scheme incurs the cost of computing <code>L0</code>, …, <code>L2</code> (which you have to do).  It also incurs</p>
<p>the cost of constructing the computation graph for <code>L0</code>, …, <code>L2</code> whether or not you end up</p>
<p>using it.  If you usually end up backpropagating through <code>L0</code>, …, <code>L2</code>, this should be okay.</p>
<p>If you usually <em>don’t</em> end up computing grads for <code>L0</code>, …, <code>L2</code>, you could:</p>
<pre><code class="lang-auto">with torch.no_grad():
    L2_out_no_graph = L2 (L1 (L0 (input)))
L2_out_no_graph.requires_grad_ (True)
output = L5 (L4 (L3 (L2_out_no_graph)))
loss = loss_fn (output)
compute_grads = should_i_compute_grads_for_l012 (output)
loss.backward()    # populates grads for L3, L4, and L5, and L2_out_no_graph
if  compute_grads:
    L2_out = L2 (L1 (L0 (input)))   # build computation graph for L0, ..., L2 at the cost of a second forward pass
    torch.autograd.backward (L2_out, grad_tensors = L2_out_no_graph.grad)   # populates grads for L0, L1, and L2


</code></pre>
<p>This scheme avoids the cost of constructing the <code>L0</code>, …, <code>L2</code> computation graph when it isn’t</p>
<p>needed (which we are assuming here is the usual case), but when you need that computation</p>
<p>graph, you not only incur the cost of building the graph, but also incur the duplicative cost of</p>
<p>recomputing <code>L2 (L1 (L0 (input)))</code>.</p>
<p>Since you don’t know in advance – that is, before computing <code>output</code> – whether you will need</p>
<p>the <code>L0</code>, …, <code>L2</code> computation graph, you can’t decide in advance which scheme to use.  You</p>
<p>have to just pick one or the other (in advance) based on the likelihood that you will end up</p>
<p>needing the <code>l0</code>, …, <code>l2</code> computation graph.</p>
<p>Best.</p>
<p>K. Frank</p>