Torchscript model is slow

jackeown: <p>I’ve made a pytorch model which uses a bit of dynamic python such as for looping where the number of loops is specified by entries of a tensor. (It’s similar to a graph neural net. I have since seen pytorch geometric, but it seems that those models can’t be exported to libtorch.) I believe that it is because of this that the model is very slow to train and run. I need to be able to export and use this model in C. I used torch.jit.script and successfully saved the model. I used libtorch to get it to C++ and a header file to make it so I can call the model from C.</p>
<p>The model is very slow to train in python. It was my understanding that a model should become much faster after running torch.jit.script on it (I thought that this turned slow python into fast C++). Would this model train faster if I trained it in C++? Why or why not? What if I just saved the model and reloaded it back into python?</p>
<p>I haven’t included the code here because I couldn’t find an easy way to reduce it to a minimal example that would be removed from the confusing specifics of my research domain. Sorry if this question is hard to grok.</p>
<p>I’m having a hard time understanding why pytorch models are slow or fast and how, in general, they can be made faster. If torchscript makes a dynamic graph static and able to run from C++, then why doesn’t it make a model much faster? (Or does it and I’m using it wrong?)</p>
<p>Thanks in advance,<br>
Jack</p>

ptrblck: <aside class="quote no-group" data-username="jackeown" data-post="1" data-topic="76389">
<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/jackeown/48/15128_2.png" class="avatar"> jackeown:</div>
<blockquote>
<p>The model is very slow to train in python. It was my understanding that a model should become much faster after running torch.jit.script on it (I thought that this turned slow python into fast C++).</p>
</blockquote>
</aside>
<p>Note that both, the Python and C++ frontend, call into the same backend methods.<br>
In the usual case you would run your model on the GPU, so that Python is used for the model “logic” and the dispatching of the kernels (to C++ and CUDA).<br>
For “big” workloads the Python overhead might be marginal.</p>
<aside class="quote no-group" data-username="jackeown" data-post="1" data-topic="76389">
<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/jackeown/48/15128_2.png" class="avatar"> jackeown:</div>
<blockquote>
<p>If torchscript makes a dynamic graph static and able to run from C++, then why doesn’t it make a model much faster?</p>
</blockquote>
</aside>
<p>If you have a static graph (with your conditions etc.), then some optimization might be applied such as fusing certain operations together.<br>
However, it depends on your actual model, if fusing ops is possible. Also note, that this area is still being worked on.</p>

jackeown: <p>Thanks for your speedy reply!</p>
<p>I understand that things like the forward method of a linear layer are using a fast backend. I’m wondering about the slow python code (like Python for-loops and other things that don’t have explicit pytorch functions).  I’m wondering if such code is made faster through torch.jit.script.</p>
<p>Perhaps my point can be seen through the code below. If I had implemented the code below in C or C++, it would be much faster than it is here. Is torchscript not able to turn these slow python for-loops into fast C for-loops? If I ran this saved module after loading it into libtorch in C++, would it run significantly faster? If so, why can’t isn’t python taking advantage of these benefits? I may be missing the point entirely, but do you have any way of making code like the following fast using pytorch?</p>
<pre><code class="lang-auto">import torch
from torch import nn
from timeit import timeit

class TestModule(nn.Module):

    def __init__(self, m, n):
        super().__init__()
        self.m = m
        self.n = n
    

    def forward(self):
        val = 0
        for i in range(self.m):
            for j in range(self.n):
                val += i*j
        return val


def purePython(m,n):
    val = 0
    for i in range(m):
        for j in range(n):
            val += i*j
    return val

m = 400
n = 400

# regular pytorch module
model = TestModule(m,n)

# jit'ed module
jitModel = torch.jit.script(model)

# saved jit'ed module
torch.jit.save(jitModel, "saved.pt")
saved = torch.jit.load("saved.pt")


numTrials = 800
print(f"pure python function time: {timeit('purePython(m,n)',  number=numTrials, globals=globals()  )}")
print(f"pytorch model time:        {timeit('model()',          number=numTrials, globals=globals()  )}")
print(f"pytorch jit model time:    {timeit('jitModel()',       number=numTrials, globals=globals()  )}")
print(f"pytorch saved model time:  {timeit('saved()',          number=numTrials, globals=globals()  )}")

</code></pre>
<p>On my computer this was the output:</p>
<pre><code class="lang-auto">pure python function time: 7.517180057002406
pytorch model time:        7.449000566004543
pytorch jit model time:    13.549996405999991
pytorch saved model time:  14.33033561299817
</code></pre>
<p>I wrote the same code in C++ and it ran in less than 1 second.<br>
I also tried importing the pytorch model in libtorch C++ and it took 13 seconds just like in python.<br>
Is there any way to make for loops like this faster in pytorch?</p>
<p>Thanks,<br>
Jack</p>

ptrblck: <p>I’m not sure, what kind of optimizations are currently implemented regarding these plain Python loops in the JIT. <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>
<p>Note that you are currently not using any tensors, but Python literals.<br>
I <em>assume</em> the first optimizations would go into the usage of tensors.<br>
If you define <code>val = torch.zeros(1)</code>, you should see a speedup of the jitted models in comparison to the eager model.</p>
<p>However, you will not see any speedup compared to the pure Python function, since your code is not vectorized.<br>
Vectorization usually gives you a good speedup and loops should be generally avoided, if possible.</p>