Keras:使用的实际GPU内存量

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 profiling工具