Neat way of temporarily disabling grads for a model?

Sam_Pr: <p>I’m implementing a version of DDPG, and trying to calculate the policy loss.</p>
<pre data-code-wrap="Python3"><code class="lang-Python3"># TODO: Should find grad for states, policy_actions but treat q_model params as constant
policy_loss = -self.q_model(states, policy_actions).mean()
</code></pre>
<p>I can do this by looping through <code>q_model</code> params, storing current values, setting <code>requires_grad</code> to <code>False</code>, calculating <code>policy_loss</code>, then restoring after.</p>
<pre data-code-wrap="Python3"><code class="lang-Python3">orig_requires_grads = [p.requires_grad for p in self.q_model.parameters()]
for p in self.q_model.parameters(): p.requires_grad = False
policy_loss = -self.q_model(states, policy_actions).mean()
for p, rg in zip(self.q_model.parameters(), orig_requires_grads): p.requires_grad = rg
</code></pre>
<p>This feels clumsy. Is there a better way to temporarily disable a model like this? I looked for something like <code>self.q_model.detach()(states, policy_actions).mean()</code> but was surprised to find it doesn’t exist. Obviously I could make a helper.</p>
<p>I plan on moving this model over to TorchScript at some point, and I’m not sure whether this is going to work out there?</p>
<p>There are some old threads that indicate that manual looping is the best way, but I want to check in case any developments have happened here.</p>
<aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/2655">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/2655" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path fill-rule="evenodd" d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/2655" target="_blank" rel="noopener nofollow ugc">[Feature request] detach (fix?) a nn.module</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2017-09-07" data-time="07:50:29" data-timezone="UTC">07:50AM - 07 Sep 17 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2020-03-12" data-time="22:28:29" data-timezone="UTC">10:28PM - 12 Mar 20 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/ruotianluo" target="_blank" rel="noopener nofollow ugc">
          <img alt="ruotianluo" src="https://avatars.githubusercontent.com/u/16023153?v=4" class="onebox-avatar-inline" width="20" height="20">
          ruotianluo
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triage review
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">Suppose, we have a siamese network.
```
out1 = net(in1)
out2 = net(in2)
```
<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">

However, we want the gradient of parameters in `net` only affected by the first branch (note that suppose we still want the gradient of `in1` and `in2`),
then what we should do here is:
```
out1 = net(in1)
for p in net.parameters(): p.requires_grad=False
out2 =net(in2)
for p in net.parameters(): p.requires_grad=True
fn(out1, out2).backward()
```


So, can we have a `detach` function for nn.module:

So, when we don't want to update a net, we can do:
`net.detach_()`
(Here, detach_() just set all the parameters requires_grad to be False)

If we want to do the example I mention above, we could do:
```
out1 = net(in1)
out2 = net.detach()(in2)
fn(out1, out2).backward()
```
(here, detach() is a clone of the input network but with all parameters requires_grad to be False)

(I give the name detach because it's essentially detach the module Parameters)</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<aside class="quote" data-post="1" data-topic="705">
  <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/alexis-jacq/48/180_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/detach-the-loaded-model/705">Detach the loaded model</a> 
  </div>
  <blockquote>
    Is there a correct way to detach a loaded model, when we don’t need to compute its gradients ? 
For the moment, I’m doing that: 
model = torch.load('mymodel.pth')
for variable in model.parameters():
    variable.detach_()

Here, I am lucky because the model contains Variable parameters, but it could contain sub-models…
  </blockquote>
</aside>


albanD: <p>Hi,</p>
<p>I don’t think there is any update. The for loop is simple and is the most efficient thing that can be done here.<br>
Especially with your special logic of things already not requiring gradients, that would be tricky.</p>
<p>Note that you can add a method to your q_model module yourself to do that to make it a bit cleaner.</p>

Sam_Pr: <p>Thanks for letting me know. Yeah, what I’ve got for now is a <code>with_grad</code> param in the forward method of my <code>q_model</code> that does what I showed. I wasn’t sure if there was a way to prevent the grad propagating through the graph without setting state like that, but it seems like there isn’t so this will have to do. I would still prefer a declarative rather than procedural API here, but it’s obviously not important.</p>

albanD: <p>Well if you don’t want any gradient flowing, you can use <code>torch.no_grad()</code>. But that means that things above won’t get gradients either.</p>

Sam_Pr: <p>In this case, the arguments given to the model should have their grads updated but the model’s own parameters shouldn’t. None of the  <code>no_grad</code>  systems allow for that specificity do they?</p>
<p>To explain the context, in DDPG we have a Q model which learns predicted rewards for state/action pairs, and a policy model which learns to choose the policy for a state. The Q model is trained from observed rewards. The policy model is trained by feeding its output into the Q model and doing gradient ascent to tweak the policy in the direction of a larger predicted reward. These both happen side-by-side in the same training loop.</p>
<p>The tricky bit is feeding the policy actions into the Q model to do the gradient ascent without also training the Q model. The  <code>requires_grad</code>  method works, but doing it in a more declarative way would require some way to tell it to only calculate grads back through a specific tensor or something like that. Like <code>with torch.no_grad(except = policy_actions):</code></p>

Sam_Pr: <p>Oops. It doesn’t particularly matter for the discussion, but there’s a little bug in the comment in my first post. It should only pass the grads back through <code>policy_actions</code>, not through <code>states</code> too. That’s a mistake that slipped in when simplifying the code to post on here, but apparently the forum won’t let me edit my first post. My real code does <code>states.detach()</code> there.</p>

albanD: <aside class="quote no-group" data-username="Sam_Pr" data-post="5" data-topic="80775">
<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/sam_pr/48/12454_2.png" class="avatar"> Sam_Pr:</div>
<blockquote>
<p>None of the <code>no_grad</code> systems allow for that specificity do they?</p>
</blockquote>
</aside>
<p>No they don’t. I just wanted to mention it in case future reader are in a different scenario.</p>

Deepesh7: <p>I am new to pytorch.<br>
So could anyone please help and clarify, if the code below would pause the computation graph at line 2(‘for p in self.q_model.parameters(): p.requires_grad = False’) and resume the graph after line 4(for p, rg in zip(self.q_model.parameters(), orig_requires_grads): p.requires_grad = rg)?</p>
<aside class="quote no-group" data-username="Sam_Pr" data-post="1" data-topic="80775">
<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/sam_pr/48/12454_2.png" class="avatar"> Sam_Pr:</div>
<blockquote>
<pre><code class="lang-auto">orig_requires_grads = [p.requires_grad for p in self.q_model.parameters()]
for p in self.q_model.parameters(): p.requires_grad = False
policy_loss = -self.q_model(states, policy_actions).mean()
for p, rg in zip(self.q_model.parameters(), orig_requires_grads): p.requires_grad = rg
</code></pre>
</blockquote>
</aside>
<p>I am trying to implement a custom loss function, where I need to pre-process output of network which involves converting tensor to numpy and then converting it back to a tensor, then finally feed it to a log loss function. Will using the code above while pre-processing, help compute gradients correctly after the pre-processing part?</p>