Backward pass with torchscript

username: Zarko
cooked: <p>Hi everyone !</p>
<p>I’m currently learning on how to use torchscript (and PyTorch in general), I have done a benchmark (CPU &amp; GPU) on differents models : gru (native PyTorch implementation), lstm (native PyTorch implementation), custom model implementation (that we will call custom_model) and jit.script(custom_model).</p>
<p>In inference I have seen that my custom model was very very slow (it’s a sort of custom GRU so that’s “normal” that he is very slow because of the loop on time step) and with the jit.script implementation my very very slow model was very good ! He reached the cudaNN implementation.</p>
<p>But I was wondering something, in training he was better than my custom implementation but very slow in comparison to the gru/lstm cudaNN implementation.</p>
<p>So my question is the following : <strong>what is in torchscript that is not doing “good” in optimizing the backward pass ?</strong></p>
<p>I have read a lot’s about this subject and the “answer” that I have found is that broadcasting is complicating the job of AD. But I haven’t fully understand well why. I see improvment so there is optimization that is doing good but not fully that would lead to the same perf as cudaNN implementation.</p>
<p>My second question is : <strong>do numba or other framework could help to reach performances of cudaNN implementation ? In my case it’s the for loop on time step (seq_len) that is generating this slower results</strong></p>
<p>custom_model = ligru<br>
jit.script(custom_model) = torch.jit.script(ligru)</p>
<p>Results are the following :<br>
<strong>Training results :</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/9/39657f9bba0babf7cc1c79d93cc9633fc5d776a9.png" data-download-href="https://discuss.pytorch.org/uploads/default/39657f9bba0babf7cc1c79d93cc9633fc5d776a9" title="timetraining"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/9/39657f9bba0babf7cc1c79d93cc9633fc5d776a9_2_690x472.png" alt="timetraining" data-base62-sha1="8bKFRNdtyQE5GR5Mi7rDkemRdwZ" width="690" height="472" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/9/39657f9bba0babf7cc1c79d93cc9633fc5d776a9_2_690x472.png, https://discuss.pytorch.org/uploads/default/original/3X/3/9/39657f9bba0babf7cc1c79d93cc9633fc5d776a9.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/9/39657f9bba0babf7cc1c79d93cc9633fc5d776a9.png 2x" data-dominant-color="EDF1F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">timetraining</span><span class="informations">766×525 34.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Thanks you a lot !</p>
<p>Take care of you guys.</p>
<p>The ressources that I have used (and others but I’m limited with links):</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/#writing-custom-rnns">
  <header class="source">
      <img src="https://pytorch.org/favicon.ico?" class="site-icon" width="32" height="32">

      <a href="https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/#writing-custom-rnns" target="_blank" rel="noopener nofollow ugc">PyTorch</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/362;"><img src="https://pytorch.org/assets/images/social-share.jpg" class="thumbnail" width="690" height="362"></div>

<h3><a href="https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/#writing-custom-rnns" target="_blank" rel="noopener nofollow ugc">Optimizing CUDA Recurrent Neural Networks with TorchScript</a></h3>

  <p>This week, we officially released PyTorch 1.1, a large feature update to PyTorch 1.0. One of the new features we’ve added is better support for fast, custom Recurrent Neural Networks (fastrnns) with TorchScript (the PyTorch JIT)...</p>


  </article>

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

  <div style="clear: both"></div>
</aside>
<p><a href="http://lernapparat.de/fast-lstm-pytorch/" class="onebox" target="_blank" rel="noopener nofollow ugc">http://lernapparat.de/fast-lstm-pytorch/</a></p>

username: Zarko
cooked: <p><strong>Inference results :</strong></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/7/172ebe6a3374dcd861d1360919ef716da9f58696.png" data-download-href="https://discuss.pytorch.org/uploads/default/172ebe6a3374dcd861d1360919ef716da9f58696" title="timeInference"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/7/172ebe6a3374dcd861d1360919ef716da9f58696_2_690x472.png" alt="timeInference" data-base62-sha1="3j583ANlzMO1m3Nab1JaxQEizT8" width="690" height="472" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/7/172ebe6a3374dcd861d1360919ef716da9f58696_2_690x472.png, https://discuss.pytorch.org/uploads/default/original/3X/1/7/172ebe6a3374dcd861d1360919ef716da9f58696.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/7/172ebe6a3374dcd861d1360919ef716da9f58696.png 2x" data-dominant-color="EDF1F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">timeInference</span><span class="informations">766×525 36.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>(I double post because I can’t post the other graphic in the main topic because I’m “new” on this forum <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> )</p>