Training speed drops a lot after doing augmentation on GPU

Zihan_LI: <p>Hi! I am doing researching on SimCLR style self-supervised learning, which requires image augmentation for contrastive learning.</p>
<ul>
<li>Initially, I pass the augmentation callable as the <code>transform</code> parameter for <code>torchvision.dataset.CIFAR10</code> constructor, which does the augmentation on cpu. I found data loading is the bottleneck for training speed.</li>
<li>I then noticed that after <code>torchvision 0.8.0</code>, we are able to integrate augmentation function, inheriting <code>nn.Module</code> , as part of the cuda routine.     This gives me around 50% speed up. I did</li>
</ul>
<pre><code class="lang-auto">from torchvision import transforms as T
transforms = nn.Sequential(
      T.RandomCrop(224),
      ...
)
</code></pre>
<ul>
<li>I kept investigating this a little bit more. Some weird things happen.<br>
– If I only do augmentation, it runs in less than 1 minute<br>
– If I only train the model (doubling the batchsize, like contrasting), it runs in less than 1 minute<br>
– But If I do augmentation then train the model, it immediately drops to 10 minute<br>
Implementation of the augmentation callable shown below and snapshots<br>
(sorry for putting everything in one pic, I’m new here. I am not allowed to upload multiple images. lol)<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/1/7171528735860726d6c01299e9fa967da61ed509.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7171528735860726d6c01299e9fa967da61ed509" title="int"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/7171528735860726d6c01299e9fa967da61ed509_2_569x500.jpeg" alt="int" data-base62-sha1="gbyKgvGIh8H9Q6sG2vaxFyxRGmB" width="569" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/7171528735860726d6c01299e9fa967da61ed509_2_569x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/7171528735860726d6c01299e9fa967da61ed509_2_853x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/1/7171528735860726d6c01299e9fa967da61ed509_2_1138x1000.jpeg 2x" data-dominant-color="F1F1F1"><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">int</span><span class="informations">1643×1442 305 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></li>
</ul>
<p>I’ve done some preliminary check, torch profiling. Everything looks good. I profiled train after augmentation for one batch, it takes 300ms to finish</p>
<p>I’ve also checked that it is not a tqdm estimating issue by training through a full epoch. Any hints on why this might happen. I run those codes on colab GPU runtime (Tesla T4)</p>

eqy: <p>How is the timing done per iteration? I’m curious if there are <code>torch.cuda.synchronize()</code> calls before timing is started and before timing is stopped as otherwise the measurements could be incorrect.</p>
<p>Additionally, is/are the size(s) of your input(s) comparable with and without augmentation?</p>