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

Question

我正在运行带有 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.08s 来实际执行推理。

但是，我注意到当我这样做时，根据nvidia-smi，我的 GPU 内存使用量激增至15637MiB / 16280MiB。

我found，您可以使用选项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.08s 才能运行。

最后，我在下面这样做了，我使用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 次以并行执行推理，这会是个问题吗？

Answer 1

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

表示使用将默认图def提交给tensorflow运行时，运行时相应地分配GPU内存。

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

如果你有 10 个会话，每个会话需要不到 1/10 的 GPU 内存，它应该可以工作。