在加载数据之前保留gpu内存张量流

Question

有没有办法在构建图形并在tensorflow中运行之前保留GPU内存？

在创建图表并运行图表之前，我需要进行缓慢的数据加载。 我正在共享一个GPU集群，所以有时其他人会在我加载数据之前使用我指定的GPU，程序会遇到OOM错误。

Answer 1

实例化tf.Session时保留设备内存。

如果您需要尽早保留设备内存，则可以先创建会话，然后继续正常构建图形。 尽管在“原始”tensorflow中首先创建图形然后实例化一个会话是常见的，但它仍然允许以不同的方式进行，并且它实际上在Keras世界中常见，其中人们经常会创建一个脚本

import tensorflow as tf
sess = tf.Session()

from keras import backend as K
K.set_session(sess)
...