JIT does not support parameter.requires_grad?

tjoseph: <p>Hi,<br>
is it a known limitation that jit.trace will ignore temporary <code>requires_grad = False</code>?</p>
<p>Here is an example:</p>
<pre><code class="lang-auto"># EXAMPLE 1
import torch
from torch import nn, jit
from torch.optim import SGD


inputs = torch.tensor([2.0], device="cuda")
model = nn.Linear(1, 1, bias=False).to("cuda")

optimizer = SGD(model.parameters(), lr=1e-1)


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = model

    def forward(self, x):
        param = next(self.parameters())

        param.requires_grad = True
        x = self.model(x).mean()
        param.requires_grad = False
        return x


c = MyModule()
forward = jit.trace(c, (inputs,))
result = forward(inputs)

result.mean().backward()

optimizer.step()
optimizer.zero_grad()

print("It does not work fine!")
</code></pre>
<pre><code class="lang-auto">RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>
<p>But when I switch the <code>requires_grad</code> flag:</p>
<pre><code class="lang-auto"># EXAMPLE 2
import torch
from torch import nn, jit
from torch.optim import SGD


inputs = torch.tensor([2.0], device="cuda")
model = nn.Linear(1, 1, bias=False).to("cuda")

optimizer = SGD(model.parameters(), lr=1e-1)


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = model

    def forward(self, x):
        param = next(self.parameters())

        param.requires_grad = False # True --&gt; False
        x = self.model(x).mean()
        param.requires_grad = True # False --&gt; True
        return x


c = MyModule()
forward = jit.trace(c, (inputs,))
result = forward(inputs)

result.mean().backward()

optimizer.step()
optimizer.zero_grad()

print("It does work fine!")
</code></pre>
<pre><code class="lang-auto">It works fine!
</code></pre>
<p>However, when I run it without jit, it runs like expected which is: example 1 runs fine and example 2 fails with an error!</p>

tjoseph: <p>Additional short question: is there something like <code>torch.jit.ignore</code> for tracing?</p>

tjoseph: <p>For everyone wondering: <code>requires_grad</code> is not supposed to work. <code>trace</code> only tracks tensor operations, not attributes. See <a href="https://github.com/pytorch/pytorch/issues/53515#issuecomment-793188191" class="inline-onebox" rel="noopener nofollow ugc">[JIT] jit.trace does not support parameter.requires_grad? · Issue #53515 · pytorch/pytorch · GitHub</a></p>
<p>Anyone know a solution? This is something that is supported in tensorflow, but not in PyTorch it seems. Makes optimizing and training my model via jit very hard even though it would profit from jit.</p>

googlebot: <blockquote>
<p>Anyone know a solution?</p>
</blockquote>
<p>script or trace deeper modules or functions - e.g. compiling(tracing) “infrastructure” MyModule makes no sense. Use <span class="mention">@jit.ignore</span> to break recursive compilation.</p>
<p>you can combine scripting &amp; tracing, if you trace functions, e.g.:</p>
<pre><code class="lang-auto">def _beta_rsample_3d(a,b):
	return D.Beta(a,b).rsample()
beta_rsample_3d = jit.trace(_beta_rsample_3d, (torch.ones(1,1,1), torch.ones(1,1,1)), check_trace=False)
</code></pre>

tjoseph: <p>Thank you for answering! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Can you explain why it does not make sense? Obviously here I use it only to showcase a minimal example, but in a more complex model, these kind of modules will exist. The solution you suggested is something I really want to avoid since it will clutter my model. I would not want to change my model to support scripting/tracing. If tracing would include tensor attributes I could just pass my whole forward pass to <code>jit.trace</code> which would properly separate optimization from model semantics. <code>jit.script</code> would be fine since it can be used at least with decorators, but unfortunately the missing support of <code>torch.distributions</code> does not allow <code>jit.script</code>.</p>
<p>btw: Do you know whether torch.jit.script will track tensor attributes (i.e. requires_grad) or is that something that the jit is generally not supposed to do and it is <em>only</em> meant for deployment?</p>

googlebot: <aside class="quote no-group" data-username="tjoseph" data-post="5" data-topic="113998">
<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/tjoseph/48/11927_2.png" class="avatar"> tjoseph:</div>
<blockquote>
<p>Can you explain why it does not make sense?</p>
</blockquote>
</aside>
<p>JIT speedups mostly come from tensor operations - like fused math operations and some optimizations specific to tensors, ops or layers. Classes like above are closer to non-computational wrappers - similar to training loop code, there is not much to optimize there.</p>
<aside class="quote no-group" data-username="tjoseph" data-post="5" data-topic="113998">
<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/tjoseph/48/11927_2.png" class="avatar"> tjoseph:</div>
<blockquote>
<p>Do you know whether torch.jit.script will track tensor attributes (i.e. requires_grad) or is that something that the jit is generally not supposed to do and it is <em>only</em> meant for deployment?</p>
</blockquote>
</aside>
<p>It is read-only (and used as meta-information in compilation). Instead, you can detach() tensors, and for parameters - you’d have to set it from outside.</p>
<aside class="quote no-group" data-username="tjoseph" data-post="5" data-topic="113998">
<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/tjoseph/48/11927_2.png" class="avatar"> tjoseph:</div>
<blockquote>
<p>If tracing would include tensor attributes I could just pass my whole forward pass to <code>jit.trace</code> which would properly separate optimization from model semantics. <code>jit.script</code> would be fine since it can be used at least with decorators, but unfortunately the missing support of <code>torch.distributions</code> does not allow <code>jit.script</code>.</p>
</blockquote>
</aside>
<p>Dunno, I find jit.trace too dangerous / limiting to use on module trees. As for torch.distributions, I mostly stopped using these wrappers when coding for performance.</p>

tom: <p>Edit: Sorry, I didn’t see you already reported it as an issue and got more input there, too.<br>
Best regards</p>
<p>Thomas</p>

tjoseph: <p>Hey, thank you for updating!</p>
<p>I see what you mean by my example not making sense. I just wanted to showcase the toggling of the requires_grad flag.</p>
<p>I am wondering why <code>.detach()</code> is supported, but not <code>requires_grad</code>. To me both fullfill orthogonal ways to influence gradient backpropagation. I guess it is just a practical limitation about how tracing works?</p>
<p>Why do you regard <code>jit.trace</code> as being dangerous/limiting? I am learning a lot here. Btw if you have some ressource where I can read up on the practical use of JIT I would appreciate you pointing me there!</p>

googlebot: <p>param.requires_grad_(b) <strong>may</strong> also work, but frankly this may be not an anticipated use case. Intuitively, you own tensors created in functions, and can toggle requires_grad freely, but parameters are owned by a module, so it is not clear if toggling works cleanly from inside.</p>
<aside class="quote no-group" data-username="tjoseph" data-post="8" data-topic="113998">
<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/tjoseph/48/11927_2.png" class="avatar"> tjoseph:</div>
<blockquote>
<p>Why do you regard <code>jit.trace</code> as being dangerous/limiting?</p>
</blockquote>
</aside>
<p>as you said above, “trace only tracks tensor operations”, this limits passable python code pretty severely. “dangerous” is associated hardcoding of non-tensor values as constants, though this issues warnings.</p>
<p>Contrived example that fails (late) with jit.trace:</p>
<pre><code class="lang-auto">k = x.size(0)
ls = [x[i] for i in range(k)]
</code></pre>

tom: <aside class="quote no-group" data-username="googlebot" data-post="9" data-topic="113998">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<pre><code class="lang-auto">k = x.size(0)
ls = [x[i] for i in range(k)]
</code></pre>
</blockquote>
</aside>
<p>Yes. So while <code>requires_grad_</code> (always use the method!) could probably be supported, the list size here inherently is invisible to PyTorch in scripting.<br>
The summary I always have is (I think it’s also in ch 15 of our book):</p>
<ul>
<li>In tracing you can do anything you want, but the JIT won’t (and can’t) try to understand it all.</li>
<li>In scripting you can only do what the understands but the JIT will do it all.</li>
</ul>
<p>This is inherent and to me it means that the second part would want to be extended in scope.</p>
<p>Best regards</p>
<p>Thomas</p>

tjoseph: <p>Thank you again for the explanation <a class="mention" href="/u/tom">@tom</a> <a class="mention" href="/u/googlebot">@googlebot</a> !</p>
<p>Something like this would solve my problem just fine:</p>
<pre><code class="lang-auto">c = MyModule()
# Tell tracer to also track some parameters
forward = jit.trace(c, (inputs,), track=c.parameters())
</code></pre>
<p>To me tracing looks like the more elegant solution if the requirements for tracing are met, which they are for me except that I need to record parameters requires_grad.</p>
<p>I guess currently there is nothing much I can do except either change my code to support scripting or not use the jit.</p>

tjoseph: <aside class="quote no-group" data-username="tom" data-post="10" data-topic="113998">
<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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p><code>requires_grad_</code> (always use the method!)</p>
</blockquote>
</aside>
<p>Could you tell me why to use the method? In PyTorch docs <code>requires_grad</code> is set directly all the time.</p>

tom: <aside class="quote no-group" data-username="tjoseph" data-post="12" data-topic="113998">
<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/tjoseph/48/11927_2.png" class="avatar"> tjoseph:</div>
<blockquote>
<p>Could you tell me why to use the method? In PyTorch docs <code>requires_grad</code> is set directly all the time.</p>
</blockquote>
</aside>
<p>First, let me admit that there is a modicum of taste here.</p>
<ol>
<li>I sometimes misspell <code>requires</code>.</li>
<li>I find myself thinking of “set requires grad to true” as an operation I apply to a tensor rather than the change of a data member. In other words, setting the gradient requirement is more “typical use of an object-modifying-method” than “typical use of <code>=</code>” to me.</li>
</ol>
<p>Best regards</p>
<p>Thomas</p>

vferrer: <p>I just find the same issue while using torch.jot.script in Pytorch 1.7.1. The easiest workaround is to just <code>x.requires_grad_(True)</code> to change the value.</p>