When should I use Tracing rather than Scripting?

username: lucidfrontier45
cooked: <p>I went through the official doc of TorchScript (<a href="https://pytorch.org/docs/stable/jit.html" rel="nofollow noopener">https://pytorch.org/docs/stable/jit.html</a>) but didn’t understand clearly what is the advantage of Tracing over Scripting.</p>
<p>Af far as I understood, both <code>jit.script</code> and <code>jit.trace</code> can convert existing <code>nn.Module</code> instances into TorchScript. However Tracing cannot handle control flow such as if/for and it also requires an example input. The inability to handle control flow sounds like a huge deal breaker.</p>
<p>The only disadvantage of Scripting I noticed is that it cannot handle several builtin modules like RNN/GRU.</p>
<p>Are there any reasons to use Tracing over Scripting?</p>
<p>Thank you.</p>
<p>BTW, can LSTM be scripted? It is a little odd that RNN and GRU cannot be scripted but LSTM can.</p>

username: driazati
cooked: <p>Tracing lets you use dynamic behavior in Python since it just records tensor operations. This may work better for your use case, but as you pointed out there are some fundamental limitations such as the inability to trace control flow / Python values. They both should be pretty easy to try out on your codebase (i.e. call <code>trace()</code> or <code>script()</code> on a module). If tracing does not work out of the box it is unlikely that you can maintain the semantics of your model and get it working under tracing (i.e. if you use control flow). Scripting may require some work to get your model using only <a href="https://pytorch.org/docs/master/jit.html#torchscript-language-reference" rel="nofollow noopener">features supported by the compiler</a>, but you will probably be able to get it working with some code changes.</p>
<p><code>GRU</code> and <code>LSTM</code> can both be compiled on master/nightly, (AFAIK only <code>LSTM</code> is in the v1.2.0 release), we haven’t had many requests for <code>RNN</code> yet. If you’d like to see it (or something else in PyTorch) be script-able that <a href="https://pytorch.org/docs/master/jit.html#builtin-functions" rel="nofollow noopener">isn’t already</a>, please file an issue on GitHub.</p>

username: lucidfrontier45
cooked: <p><a class="mention" href="/u/driazati">@driazati</a> Thank you for reply.</p>
<p>Do you mean that Tracing can handle almost all Python features/libraries except for for/while/if kind flows while Scripting can only handle a subset of Python features (aka TorchScript)? I still suspect there are many Python features that cannot be used with Tracing and therefore the difference between Tracing and Scripting is very small. Could you kindly give me an example where we should use Tracing over Scripting?</p>

username: driazati
cooked: <p>Tracing can handle anything that uses only PyTorch tensors and PyTorch operations. If someone passed a PyTorch tensor to a Pandas dataframe and did some operations, tracing wouldn’t capture that (though neither would script at this point), so there are limitations. If the only data flowing around your computations are tensors and there is no control flow, tracing is probably the way to go. Otherwise, use scripting.</p>
<p>The <a href="https://github.com/facebookresearch/pytext" rel="nofollow noopener">pytext</a> library uses a mix of <a href="https://github.com/facebookresearch/pytext/blob/2db0cc9d266b4ae45c5911a530f950bb49f79273/pytext/models/doc_model.py#L137" rel="nofollow noopener">scripting</a> and <a href="https://github.com/facebookresearch/pytext/blob/739aada896883235599d2589745b635347132904/pytext/models/test/bilstm_test.py" rel="nofollow noopener">tracing</a>, and it all generally works well since they can be <a href="https://pytorch.org/docs/master/jit.html#mixing-tracing-and-scripting" rel="nofollow noopener">mixed together</a> pretty seamlessly.</p>

username: vincentzlt
cooked: <p>I have got similar confusion.</p>
<p>Seems that even If “the only data flowing around your computations are tensors and there is no control flow”, we can still use scripting. Does this mean we can actually use scripting for all cases? Any speed difference between the two?</p>
<p>Thanks.</p>

username: driazati
cooked: <p>Right, the only thing that would work in tracing but not scripting is use of Python language features and dynamic behavior that script mode doesn’t support. Since scripting compiles your code, you may have to do some work to make the compiler happy (e.g. add type annotations). But it is easy to try, just pass your module to <code>torch.jit.script</code>. The two compile to the same IR under the hood, so the speed should be about the same.</p>

username: Midhilesh
cooked: <p>Hi,<br>
I don’t understand what do you mean by dynamic behaviour and python language features that is supported in trace but not in script? also</p>
<p>what is the difference between Torchscript compiler and JIT compiler?</p>
<pre><code class="lang-auto">Scripting a function or  `nn.Module`  will inspect the source code, 
compile it as TorchScript code using the TorchScript compiler. 
</code></pre>
<pre><code class="lang-auto">Trace a function and return an executable 
that will be optimized using just-in-time compilation.
</code></pre>
<p>I request you to explain those in detail.</p>
<p>Thanks.</p>

username: austinmw
cooked: <p>Hi, would you be able to give an example of dynamic behavior / Python language features that script mode doesn’t support? It sounds a bit abstract to me compared to the control flow statements that I know only work with script mode.</p>

username: mpariente
cooked: <p>It was also abstract for me in which can trace would be superior/more useful than scripting. I finally found an example in my code</p>
<pre><code class="lang-python">import torch
from torch import nn


class MyModule(nn.Module):
    def __init__(self, return_b=False):
        super().__init__()
        self.return_b = return_b

    def forward(self, x):
        a = x + 2
        if self.return_b:
            b = x + 3
            return a, b
        return a


model = MyModule(return_b=True)

# Will work
traced = torch.jit.trace(model, (torch.randn(10, ), ))

# Will fail
scripted = torch.jit.script(model)
</code></pre>
<p>This can easily be changed to be scriptable, but if you know the control flow is static once the model is exported, tracing will work just fine.</p>

username: ppwwyyxx
cooked: <p>I wrote here <a href="https://ppwwyyxx.com/blog/2022/TorchScript-Tracing-vs-Scripting/" class="inline-onebox" rel="noopener nofollow ugc">TorchScript: Tracing vs. Scripting - Yuxin's Blog</a> on the advantage of tracing, and why it should be preferred over scripting for any large projects.</p>