Pytorch is (much) slower than Tensorflow

Eli_Osherovich: <p>I am new to Pytorch and while trying it I felt that familiar models run slower than in Tensorflow.</p>
<p>To illustrate the behavior, I built a very simple fully connected model (single layer) which is available here:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1NaSC1X7RtxNKCV4jQEJ_HwNGgVq8viPc?usp=sharing">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/b13f77f4b3ddf7f66c89357b93aa7e33/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1NaSC1X7RtxNKCV4jQEJ_HwNGgVq8viPc?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar" width="260" height="260">

<h3><a href="https://colab.research.google.com/drive/1NaSC1X7RtxNKCV4jQEJ_HwNGgVq8viPc?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

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

  <div style="clear: both"></div>
</aside>

<p>Tensorflow runs three times faster than Pytorch.</p>
<p>Is such a ratio expected? Do I  miss something in PT?</p>

ptrblck: <p>Your workload is tiny as the model contains a single small linear layer only and you might thus see the overheads of the kernel launches, the data loading etc.<br>
You could try to use <code>torch.compile</code> with <code>mode='reduce-overhead'</code>, which would internally use CUDA Graphs, but I doubt this would help since you are using a single layer in the end.</p>

Eli_Osherovich: <p>Where/when does this overhead happen? If it happens every epoch (why?) it can somehow explain it. However, it it happens only once, I hardly see how the impact can be so severe over 20 epochs…</p>
<p>I tried compilation with different flags, but it only made things (a bit) worse…</p>
<p>P. S.<br>
Is there any chance that the developers look into this and, probably, find a way to optimize PT for small problems as well?</p>

ptrblck: <aside class="quote no-group" data-username="Eli_Osherovich" data-post="3" data-topic="174519">
<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/eli_osherovich/48/58200_2.png" class="avatar"> Eli_Osherovich:</div>
<blockquote>
<p>Where/when does this overhead happen? If it happens every epoch (why?) it can somehow explain it. However, it it happens only once, I hardly see how the impact can be so severe over 20 epochs…</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="174519">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>and you might thus see the overheads of the kernel launches, the data loading etc.</p>
</blockquote>
</aside>
<p>The data loading is performed by the <code>DataLoader</code> and yields a batch in each iteration. If the data loading time exceeds the model training time (which could easily be the case for a single tiny <code>matmul</code>) you will see this overhead in each iteration.<br>
The kernel launch overheads are visible in all kernel launches which are triggered after PyTorch went through its dispatching logic. This overhead is visible the more tiny kernels you want to schedule and CUDA Graphs can help here.</p>
<aside class="quote no-group" data-username="Eli_Osherovich" data-post="3" data-topic="174519">
<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/eli_osherovich/48/58200_2.png" class="avatar"> Eli_Osherovich:</div>
<blockquote>
<p>Is there any chance that the developers look into this and, probably, find a way to optimize PT for small problems as well?</p>
</blockquote>
</aside>
<p>CUDA Graphs was added for it: <a href="https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/" class="inline-onebox">Accelerating PyTorch with CUDA Graphs | PyTorch</a></p>

Eli_Osherovich: <p>Data loading should be identical in both cases – this is an in-memory numpy array.<br>
Not sure if cuda graphs will be beneficial on CPU… Will they?</p>