Keras:实际使用的GPU内存量 [英] Keras: real amount of GPU memory used

问题描述

我将Keras与Tensorflow后端一起使用，并且查看nvidia-smi不足以了解当前网络体系结构需要多少内存，因为Tensorflow似乎只是分配了所有可用的内存.

I'm using Keras with Tensorflow backend and looking at nvidia-smi is not sufficient to understand how much memory current network architecture need because seems like Tensorflow just allocate all availible memory.

所以问题是如何找出实际的GPU内存使用情况?

So the question is how to find out real GPU memory usage?

推荐答案

可以使用时间轴来完成，它可以为您提供有关内存日志记录的完整跟踪.类似于下面的代码:

It can be done using Timeline, which can give you a full trace about memory logging. Similar to the code below:

from keras import backend as K
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完成.类似于以下代码:

If you want to limit the gpu memory usage, it can alse be done from gpu_options. Like the following code:

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))

检查有关时间轴的以下文档

在后端使用TensorFlow时，可以使用

As you use TensorFlow in the backend, you can use tfprof profiling tool

这篇关于Keras:实际使用的GPU内存量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！