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 来实际执行推断。

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

我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内存。

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

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