时间戳

首页 > 解决方案 > TensorFlow 在执行推理时如何分配 GPU 内存？

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

问题描述

我正在运行带有 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.08 秒才能运行。

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

标签： tensorflowneural-networkcomputer-vision

解决方案

让我们首先确保发生了什么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 内存，它应该可以工作。

推荐阅读