Question

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

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

Answer 1

可以使用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分析工具

Answer 2

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

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

Keras：使用的实际GPU内存量

2 个答案: