我在一个计算资源共享的环境中工作,也就是说,我们有几台服务器机器,每台机器都配备了几个Nvidia Titan X gpu。

For small to moderate size models, the 12 GB of the Titan X is usually enough for 2–3 people to run training concurrently on the same GPU. If the models are small enough that a single model does not take full advantage of all the computational units of the GPU, this can actually result in a speedup compared with running one training process after the other. Even in cases where the concurrent access to the GPU does slow down the individual training time, it is still nice to have the flexibility of having multiple users simultaneously train on the GPU.

TensorFlow的问题在于,默认情况下,它在启动时分配了全部可用的GPU内存。即使是一个小型的两层神经网络,我看到所有12 GB的GPU内存都用完了。

有没有一种方法让TensorFlow只分配,比如说,4 GB的GPU内存,如果我们知道这对一个给定的模型来说已经足够了?


当前回答

上面所有的答案都假设使用sess.run()调用来执行,这在TensorFlow的最新版本中成为异常而不是规则。

当使用tf。估计器框架(TensorFlow 1.4及以上)将分数传递给隐式创建的MonitoredTrainingSession的方式是,

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
trainingConfig = tf.estimator.RunConfig(session_config=conf, ...)
tf.estimator.Estimator(model_fn=..., 
                       config=trainingConfig)

类似地,在Eager模式下(TensorFlow 1.5及以上),

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
tfe.enable_eager_execution(config=conf)

编辑:11-04-2018 例如,如果要使用tf.contrib.gan。Train,那么你可以使用类似bellow的东西:

tf.contrib.gan.gan_train(........, config=conf)

其他回答

无耻插头:如果你安装了Tensorflow支持的GPU,无论你设置它只使用CPU还是GPU,会话都会首先分配所有的GPU。我可以补充我的提示,即使你设置图形只使用CPU,你也应该设置相同的配置(如上所述:)),以防止不必要的GPU占用。

在像IPython和Jupyter这样的交互界面中,您也应该设置该配置,否则,它将分配所有内存,而几乎没有内存留给其他内存。这一点有时很难注意到。

对于TensorFlow 2.0和2.1 (docs):

import tensorflow as tf
tf.config.gpu.set_per_process_memory_growth(True)

对于TensorFlow 2.2+ (docs):

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

文档还列出了更多的方法:

设置环境变量TF_FORCE_GPU_ALLOW_GROWTH为true。 使用tf.config.experimental。set_virtual_device_configuration设置虚拟GPU设备的硬限制。

对于Tensorflow 2.0,这个解决方案很适合我。(TF-GPU 2.0, Windows 10, GeForce RTX 2070)

physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)

如果你正在使用Tensorflow 2,请尝试以下步骤:

config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.Session(config=config)
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

https://github.com/tensorflow/tensorflow/issues/1578