List all the tensors and their memory allocation

1

Hello everyone!

Is there a way to list all the tensors and their memory usage? I run out of GPU memory when I start to infer a trained model (not training at all in this code). So I would like to know where is the bottleneck. What uses up all the memory.

Edit: My saved models are more than 700MB in size on disk

2

You could use a similar approach as described in this post to get all Python tensors.
However, this would not return the tensors allocated in the backend, so you might (additionally) want to check the memory usage via e.g. print(torch.cuda.memory_summary()) to narrow down where you are running out of memory.

2 Likes
3

Thank you so much! That is very helpful.
I could not spot anything unusual. I am running a simple Unet. I presume these networks use a lot of memory. It would have been handy if we can list tensors by their name and their memory usage (size might not tell the full story because of the underlying data type I guess). But it might not possible I think. Thanks again!