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

How to calculate GPU memory usage in Theano?

我正在试验不同的 Theano 模型,并使用序列长度不断增加的课程。我如何才能提前预测对于任何给定的序列长度和模型,批量大小要有多大才能填充 GPU 的内存?

更糟糕的是,如果我不小心使用了太多内存,我会得到一个 MemoryError 并且 GPU 上的内存没有被释放,要求我重新启动进程以释放内存,并在之前丢失我的网络尝试新的批量大小。因为这个错误是不可恢复的,所以很难只增加批量大小直到出现异常然后再退缩。

似乎 Theano 没有任何 built-in 方法来估计模型的内存大小。最好的办法是创建一个已知大小的模型子集,并使用 Theano 手册中 here 描述的内存估计技术。

我们还需要考虑我们的对象在 GPU 中的表示方式(例如,我们使用 float32 还是 float64,以及每个对象在内部占用多少字节GPU)。

一旦您可以估计一个小模型的大小,您就可以将这些估计以合理的精度投影到一个更大的模型的大小。您应该能够编写自己的内存估计函数,该函数可以将许多特征和观察值、张量或图形节点作为参数,returns 内存使用量。

假设您知道要存储在 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

假设 over-head 常量为 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

通过监控可用内存量并计算您的 model/data 大小,您可以更好地使用 GPU 内存。但是,请注意 memory fragmentation 问题,因为它可能会导致 MemoryError 意外。