How to profile memory in Pytorch

What are the standard ways of profiling memory in pytorch?

I have a model, and I want to find out where the memory is spent during training. I can iterate over gc.get_objects, but only tensors are included there. I think, most of the memory is spent on backprop data. Is there any way to find out how much memory is consumed on backprop?


nvidia-smi is your best bet (when using GPU’s). For CPU’s I’ve found useful psutil.virtual_memory().used
Also, pytorch-summary can take in a model and give estimates for the memory needed for the parameters, forward/backward pass, and the data itself.