CUDA out of memory during training

username: PauleB
cooked: <p>Hello,</p>
<p>I am pretty new to machine learning and I am facing an issue I cannot solve by myself.</p>
<p>I took this code to implement U-net model and modify it a little bit to fit my dataset: <a href="https://www.kaggle.com/hsankesara/unet-image-segmentation" rel="nofollow noopener">https://www.kaggle.com/hsankesara/unet-image-segmentation</a></p>
<p>I have 240 pairs of images-masks both 256x256 px size.<br>
My batch size is reduce to 8 but I have the error “CUDA out of memory. Tried to allocate 408.00 MiB (GPU 0;5.00 GiB total capacity; 3.00 GiB already allocated; 236.37 MiB free; 3.46 GiB reserved in total by PyTorch)”.</p>
<p>The error happens at the epoch 49/1000 in the middle of the forward function of the Unet class.</p>
<p>From suggestions to same issue topics, I have tried to del some variables and use torch.cuda.empty_cache() after each iterations, I have also tried to use torch.no.grad() to get the validation loss.<br>
Those did not help, but it is possible that I did not put them in the right place.</p>
<p>From the linked code, is it possible to have at least advice to where I should empty the cache, delete the variables etc ?<br>
Or if you have other hypothesis to solve this issue of memory ?</p>
<p>Thank you for the help !</p>
<p>Paule.</p>

username: ptrblck
cooked: <p>You don’t need to call <code>torch.cuda.empty_cache()</code>, as it will only slow down your code and will not avoid potential out of memory issues.<br>
If PyTorch runs into an OOM, it will automatically clear the cache and retry the allocation for you.</p>
<p>That being said, you shouldn’t accumulate the <code>batch_loss</code> into <code>total_loss</code> directly, since <code>batch_loss</code> is still attached to the computation graph, which will be stored as well.<br>
Use <code>total_loss += batch_loss.detach()</code> instead and rerun the code.</p>
<p>Also, you are right in wrapping the validation loop into a <code>with torch.no_grad()</code> block, as this will avoid storing the intermediate tensors, which would be needed for the backward pass.</p>
<p>If that doesn’t help, you might need to further decrease the batch size or use <code>torch.utils.checkpoint</code> to trade compute for memory.</p>

username: PauleB
cooked: <p>Hey,<br>
Thank you for the advice!<br>
I have tried to reduce the batch size to 4, the checkpoints, the detach() for the batch loss and I have also tried the torch.cuda.amp.GradScaler().<br>
But nothing changed, the error still happened at the same place <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"> which I found quite weird.<br>
I will try to investigate it further, maybe reducing the images size (but for me, if I do that it is more about putting a plaster on the problem instead of solving it).</p>

username: PauleB
cooked: <p>Update:<br>
The problem came from the validation calculation actually. I do a validation loss calculation each 50 epochs and my validation dataset size is 72 which is too large. That is why I have an error " CUDA out of memory" while finishing the epoch 49 and starting the epoch 50.<br>
I will change the code in that way and hopefully it will work !</p>

username: Bohao_Cheung
cooked: <p>A good suggestion that use with torch.no_grad() in test and validation phase(clear immediate tensors) and detach the loss(remove cached calculating phase) while calculating the total loss!!!</p>