Memory usage reduction

Adam_Gregor: <p>Hello,<br>
I am trying to perform Supervised Contrastive Learning according to <a href="https://github.com/HobbitLong/SupContrast" rel="noopener nofollow ugc">https://github.com/HobbitLong/SupContrast</a> which was created by one of co-authors of the original paper.<br>
I am trying to use a dataset of my own and input size of images 256x256px (instead of the original 32x32). Hence I set the parameter <code>size</code> of <code>main_supcon.py</code> to 256, put it my dataset and run the script. However, as soon as I start putting batches (of size 64) into the newly created network, I recieve</p>
<pre><code class="lang-auto">RuntimeError: CUDA out of memory.
</code></pre>
<p>error. This however seems suspicious to me, as I am using a cluster with 2x  Nvidia A40 GPUs with total of 90GB VRAM:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/3/039fcd9d85cccc27a5d7f27c8a8bef8b4502b041.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/039fcd9d85cccc27a5d7f27c8a8bef8b4502b041" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/3/039fcd9d85cccc27a5d7f27c8a8bef8b4502b041_2_690x251.jpeg" alt="image" data-base62-sha1="w3Oaov4HT9obMJOYEtLeSoVVe1" width="690" height="251" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/3/039fcd9d85cccc27a5d7f27c8a8bef8b4502b041_2_690x251.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/0/3/039fcd9d85cccc27a5d7f27c8a8bef8b4502b041.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/3/039fcd9d85cccc27a5d7f27c8a8bef8b4502b041.jpeg 2x" data-dominant-color="3B1730"><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">image</span><span class="informations">786×286 99.4 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> and from the source code I got the feeling that it is capable to utilize several GPUs. Nevertheless, even if only one GPU was in use, it would lead aproximately to 2x64x256x256x3 = 25165824B ~= 25. GB of VRAM.<br>
My question hence is: does anyone knows what may cause this out of memory errors and how to fix them?</p>
<p>I am using PyTorch for cuda version 11.8.</p>
<p>Thank you,<br>
Adam</p>

ptrblck: <aside class="quote no-group" data-username="Adam_Gregor" data-post="1" data-topic="175236">
<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/adam_gregor/48/44783_2.png" class="avatar"> Adam_Gregor:</div>
<blockquote>
<p>and from the source code I got the feeling that it is capable to utilize several GPUs.</p>
</blockquote>
</aside>
<p>Could you verify if this is indeed the case and if e.g. <code>DistributedDataParallel</code> is used or any model sharding approach?</p>
<aside class="quote no-group" data-username="Adam_Gregor" data-post="1" data-topic="175236">
<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/adam_gregor/48/44783_2.png" class="avatar"> Adam_Gregor:</div>
<blockquote>
<p>Nevertheless, even if only one GPU was in use, it would lead aproximately to 2x64x256x256x3 = 25165824B ~= 25. GB of VRAM.</p>
</blockquote>
</aside>
<p><code>25165824</code> bytes will use ~25MB not GB. Assuming this shape represents the model input, you would have to multiple it by <code>4</code> since each value would use <code>float32</code> in the default setup (4 bytes per value).<br>
However, the input is usually using the least amount of memory while the model parameters and the intermediate forward activations, needed for the gradient calculation, might use significantly more memory. Also, the gradients as well as optimizer states (if applicable) will increase the memory usage.</p>