GPU RAM Profiling

Any tips on profiling gpu memory over time? Specifically, I would like to see a breakdown of which objects are taking up memory and how much each one is consuming. Even seeing that breakdown during peak consumption would be very helpful.

I was going through the image segmentation assignment(lesson 3 camvid) and I noticed that I couldn’t use a batch size larger than 1 without running out of gpu memory(8gb gtx 1080). I’m totally prepared to accept that my card is too small, except where is all of that memory going? In terms of what needs to be loaded onto my GPU during training it’s just the model, gradients for each weight and a training batch, right?

For a similar pix2pix problem using a unet, a unet model saved to disk is ~500 mb. I’m assuming that it would be roughly the same size on gpu ram and that the gradients would also be the same size.

Unless I’m missing something:

consumption = model + batch + grads
model = 500 mb
grads = bs * model
batch = bs * num_channels * width * height * 2 (for inputs and outputs) * wordsize

The images are 600x800x3 and the word size is 4 bytes(32 bit float). This means
batch = bs * 600 * 800 * 3 * 2 *4 = bs * 11 mb

then consumption = 500mb + (511mb) * bs.

If bs=4, I should only be consuming ~2.5 gb right? Where is the rest of the memory going?