如何计算Theano中的GPU内存使用情况？

我正在尝试不同的Theano模型,并使用不断增加的序列长度的课程.为了填充GPU的内存,我怎样才能提前预测任何给定序列长度和模型的批量大小有多大？

更糟糕的是,如果我不小心使用了太多内存,我会得到一个MemoryError并且GPU上的内存没有释放,要求我重新启动进程以释放内存,并在尝试新批量之前丢失网络.由于此错误无法恢复,因此很难在异常之前增加批量大小,然后再向下调整.

假设您知道要存储在GPU上的元素数量,您可以轻松计算存储这些元素所需的内存量.

一个简单的例子:

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

假设过头常数为0将打印:

>>> Data will need 1.22 GBs of free memory

如果您使用的是NVIDIA显卡并在计算机上安装了CUDA,那么您可以使用以下代码行轻松获取GPU上的总可用内存量:

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))

然后输出采用以下格式(对于我的机器):

>>> 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

通过监控可用内存量并计算模型/数据的大小,您可以更好地使用GPU内存.但是,请注意内存碎片问题,因为它可能会导致MemoryError意外.