当执行推理时，TensorFlow如何分配GPU内存？

Question

当执行推理时，TensorFlow如何分配GPU内存？

4

我正在使用ResNet50架构运行FastRCNN。我通过以下方式加载模型检查点并进行推理：

saver = tf.train.Saver()
saver.restore(sess, 'model/model.ckpt')
with tf.Session() as sess:
    sess.run(y_pred, feed_dict={x: input_data})

一切看起来都运行得很顺利。该模型需要0.08秒才能真正执行推论。

但是，我注意到当我这样做时，我的GPU内存使用量会爆炸到15637MiB / 16280MiB，这是根据nvidia-smi显示的。

我发现可以使用选项config.gpu_options.allow_growth来阻止Tensorflow分配整个GPU，并改为按需使用GPU内存：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True

saver = tf.train.Saver()
saver.restore(sess, 'model/model.ckpt')
with tf.Session(config=config) as sess:
    sess.run(y_pred, feed_dict={x: input_data})

这样做可以将内存使用减少到4875MiB / 16280MiB。模型仍然需要0.08秒运行。

最后，我采用了下面的方法，使用per_process_gpu_memory_fraction为模型分配固定数量的内存。

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.05

saver = tf.train.Saver()
saver.restore(sess, 'model/model.ckpt')
with tf.Session(config=config) as sess:
    sess.run(y_pred, feed_dict={x: input_data})

将此操作执行后，使用量降至1331MiB / 16280MiB，模型仍需要0.08s才能运行。

这引出一个问题 - TF在推理时如何为模型分配内存？如果我想在同一GPU上加载此模型10次以并行执行推理，是否会有问题？

- farza

一如既往。如果不使用allow growth，它只会分配所有内存，并使用与您的模型权重相同的内存量。直到您的数据可以适应内存为止，它才不会引起问题。 - Sharky

2

那么，当我使用 per_process_gpu_memory_fraction 将 GPU 内存降至 1331MiB / 16280MiB 时，为什么一切仍然能够正常运行，并且与使用 allow_growth 时完全相同？当使用 allow_growth 时，内存为 4875MiB / 16280MiB。 - farza

这里描述得非常好：https://www.tensorflow.org/guide/using_gpu - Sharky

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- pinxue · Accepted Answer

首先让我们确保在tf.Session(config=config)中发生了什么。

这意味着将默认的图形定义提交给TensorFlow运行时，运行时会相应地分配GPU内存。

然后Tensorflow将分配所有GPU内存，除非您通过设置per_process_gpu_memory_fraction来限制它。如果无法分配所需的内存量，则会失败，除非.gpu_options.allow_growth = True，这告诉TF在失败的情况下尝试重新分配较少的内存，但迭代始终从全部或部分GPU内存开始。

如果您有10个会话，每个会话需要少于1/10的GPU内存，则应该可以正常工作。