How to calculate GPU memory usage in Theano?

Question

I am experimenting with different Theano models and use a curriculum of ever increasing sequence length. How can I predict ahead of time how big to make the batch size for any given sequence length and model in order to fill the GPU's memory?

To make matters worse, if I ever accidentally use too much memory, I get a MemoryError and the memory on the GPU is not freed, requiring me to restart the process to free the memory, and lose my network, before trying a new batch size. Because this error is unrecoverable, it is difficult to just increase batch size until exception and then back down.

Answer 1

Assuming you know the number of elements to be stored on a GPU you can easily compute the amount of memory required to store those elements.

A simple example:

import numpy as np
import theano.tensor as T
T.config.floatX = 'float32'
dataPoints = np.random.random((5000, 256 * 256)).astype(T.config.floatX)
#float32 data type requires 4 bytes
sizeinGBs = 5000 * 256 * 256 * 4 / 1024. / 1024 / 1024 + (some small over-head constant)
print "Data will need %2f GBs of free memory" % sizeInGB

Assuming the over-head constant is 0 will print:

>>> Data will need 1.22 GBs of free memory

If you are using an NVIDIA graphics card and have installed CUDA on your machine then you can easily get the total amount of free memory on your GPU using the following line of code:

import theano.sandbox.cuda.basic_ops as sbcuda
import numpy as np
import theano.tensor as T
T.config.floatX = 'float32'
GPUFreeMemoryInBytes = sbcuda.cuda_ndarray.cuda_ndarray.mem_info()[0]
freeGPUMemInGBs = GPUFreeMemoryInBytes/1024./1024/1024
print "Your GPU has %s GBs of free memory" % str(freeGPUMemInGBs)
#An operation is to be executed below
testData = shared(np.random.random((5000, 256 * 256)).astype(T.config.floatX), borrow = True)
print "The tasks above used %s GBs of your GPU memory. The available memory is %s GBs" % (str(freeGPUMemInGBs - GPUFreeMemoryInBytes/1024./1024/1024), str(GPUFreeMemoryInBytes/1024./1024/1024))

Then the output is in the following format (for my machine here):

>>> Your GPU has 11.2557678223 GBs of free memory
>>> The tasks above used 1.22077941895 GBs of your GPU memory. The available memory is 10.0349884033 GBs

By monitoring the amount of free memory and calculating the size of your model/data you can better use the GPU memory. However, be aware of the memory fragmentation issue as it may cause MemoryError unexpectedly.