How to calculate GPU memory usage in Theano?

Question

I experiment with various Theano models and use a curriculum with an ever-increasing sequence length. How can I predict in advance how large the lot size is for any length and sequence model to fill the GPU memory?

To make matters worse, if I ever accidentally used too much memory, I get a MemoryError and the memory on the GPU is not freed, requiring a process restart to free memory and network loss before trying a new batch size. Since this error cannot be fixed, it is difficult to simply increase the batch size to an exception, and then return down.

Answer 1

Assuming that you know the number of elements that should be stored on the GPU, you can easily calculate the amount of memory needed to store these 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 constant above the head is 0, it will print:

 >>> Data will need 1.22 GBs of free memory 

If you use an NVIDIA graphics card and installed CUDA on your computer, 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 controlling 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 problem of memory fragmentation , as this may cause a MemoryError unexpectedly.

Answer 2

Theano doesn't seem to have a built-in way to estimate the memory size of the model. It is best to create a small subset of your model with a known size and use the memory estimation methods described here in Theano's manual.

We also need to consider how our objects are represented inside the GPU (for example, we use float32 or float64 , and how many bytes each of them occupies in 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, which can take a number of functions and observations, or tensors, or graph nodes as parameters and returns the amount of memory used.