Keras：实际使用的 GPU 内存量答案

【问题标题】：Keras: real amount of GPU memory usedKeras：实际使用的 GPU 内存量
【发布时间】：2017-05-18 14:31:31
【问题描述】：

我正在使用带有 Tensorflow 后端的 Keras，而查看 nvidia-smi 不足以了解当前网络架构需要多少内存，因为 Tensorflow 似乎只是分配了所有可用内存。

所以问题是如何找出真正的 GPU 内存使用情况？

【问题讨论】：

你试过'model.summary()'吗？它应该对模型内存使用情况有所了解。
@orabis 是的，但这只是权重，如果我们训练模型 Tensorflow 还会分配迭代 blob 和梯度 blob + 一些开销，我不知道如何精确计算内存使用量。

【解决方案1】：

可以使用 Timeline 来完成，它可以为您提供有关内存记录的完整跟踪。类似于下面的代码：

from keras import backend as K
from tensorflow.python.client import timeline
import tensorflow as tf


with K.get_session()  as s:
    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
     
    # your fitting code and s run with run_options 

    to = timeline.Timeline(run_metadata.step_stats)
    trace = to.generate_chrome_trace_format()
    with open('full_trace.json', 'w') as out:
            out.write(trace)

如果你想限制 gpu 内存使用，也可以通过 gpu_options 来完成。像下面的代码：

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.2
set_session(tf.Session(config=config))

查看以下有关 Timeline 对象的文档

当您在后端使用 TensorFlow 时，您可以使用tfprof 分析工具

【讨论】：

问题是如果你运行fit，然后是session.run，当你的检测run调用开始时，fit分配的大部分内存将被释放。这里有一个相关的问题：github.com/tensorflow/tensorflow/issues/9868 缺少的是让 Keras 使用自定义 run_options 或向 keras 模型添加自定义操作（如 MaxBytesInUse op 像 here ）的配方

【解决方案2】：

在告诉 TensorFlow 不要保留 GPU 的所有内存，而是按需增加此保留后，您仍然可以使用 nvidia-smi：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
keras.backend.tensorflow_backend.set_session(tf.Session(config=config))

【讨论】：