如何衡量 TensorFlow 模型的 GPU 内存使用情况

Question

我的目标是弄清楚在推理过程中保存为 .pb 文件的 TensorFlow 模型使用了多少 GPU 内存。 通过将per_process_gpu_memory_fraction限制为 0.01 的值并在输入图像上成功运行模型，我预计内存使用量为 120MB（基于 12,000MB GPU）。

conf = tf.ConfigProto()
conf.gpu_options.per_process_gpu_memory_fraction=0.01
session = tf.Session(config=conf)

当我通过nvidia-smi观察实际的 GPU 内存使用情况时，我看到了 550MB 的利用率。 根据此处的答案 ( https://stackoverflow.com/a/55532954 )，我测量了创建 TF 会话所需的默认内存，总计 150MB。 现在的问题是剩余的 550MB（测量值）- 150MB（会话）- 120MB（模型）= 280MB 来自哪里。 我想排除使用 GPU 内存的任何其他因素，只量化模型本身使用了多少。 在这种情况下，我可以简单地假设 120MB 吗？

Answer 1

在运行 python 代码之前，尝试将 TF_FORCE_GPU_ALLOW_GROWTH 设置为 true：

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

在导入 tensorflow 之前放置它。

然后测量 GPU 内存利用率。 此外，它还取决于如何分配 GPU 上的内存的 Tensorflow 内存分配器。 所以它通常会比内存中的实际模型大一点。

还可以尝试这里提到的建议： 如何限制 tensorflow GPU 内存使用？