I have 240 pairs of images-masks both 256x256 px size.
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)”.
The error happens at the epoch 49/1000 in the middle of the forward function of the Unet class.
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.
Those did not help, but it is possible that I did not put them in the right place.
From the linked code, is it possible to have at least advice to where I should empty the cache, delete the variables etc ?
Or if you have other hypothesis to solve this issue of memory ?
You don’t need to call torch.cuda.empty_cache(), as it will only slow down your code and will not avoid potential out of memory issues.
If PyTorch runs into an OOM, it will automatically clear the cache and retry the allocation for you.
That being said, you shouldn’t accumulate the batch_loss into total_loss directly, since batch_loss is still attached to the computation graph, which will be stored as well.
Use total_loss += batch_loss.detach() instead and rerun the code.
Also, you are right in wrapping the validation loop into a with torch.no_grad() block, as this will avoid storing the intermediate tensors, which would be needed for the backward pass.
If that doesn’t help, you might need to further decrease the batch size or use torch.utils.checkpoint to trade compute for memory.
Hey,
Thank you for the advice!
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().
But nothing changed, the error still happened at the same place which I found quite weird.
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).
Update:
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.
I will change the code in that way and hopefully it will work !
which I found quite weird.