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.

Answer 2

It seems that Theano doesn't have any built-in way to estimate the memory size of a model. Your best bet is to create a small subset of your model with a known size, and use the memory estimation techniques described here in the Theano manual.

We'll also need to take into account how our objects are represented inside the GPU (are we using float32 or float64, for example, and how many bytes does each take up inside the GPU).

Once you can estimate the size of a small model, you can then project these estimates to the size of a much larger model with reasonable accuracy. You should be able to write your own memory estimation function that can take a number of features and observations, or tensors, or graph nodes as parameters and returns a memory usage amount.