我在一个计算资源共享的环境中工作,也就是说,我们有几台服务器机器,每台机器都配备了几个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内存,如果我们知道这对一个给定的模型来说已经足够了?


当前回答

Tensorflow 2.0 Beta和(可能)更高版本

API再次改变。现在可以在以下地方找到它:

tf.config.experimental.set_memory_growth(
    device,
    enable
)

别名:

tf.compat.v1.config.experimental.set_memory_growth tf.compat.v2.config.experimental.set_memory_growth

引用:

https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/config/experimental/set_memory_growth https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth

参见: Tensorflow—使用GPU: https://www.tensorflow.org/guide/gpu

对于Tensorflow 2.0 Alpha,请参见:这个答案

其他回答

Tensorflow 2.0 Beta和(可能)更高版本

API再次改变。现在可以在以下地方找到它:

tf.config.experimental.set_memory_growth(
    device,
    enable
)

别名:

tf.compat.v1.config.experimental.set_memory_growth tf.compat.v2.config.experimental.set_memory_growth

引用:

https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/config/experimental/set_memory_growth https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth

参见: Tensorflow—使用GPU: https://www.tensorflow.org/guide/gpu

对于Tensorflow 2.0 Alpha,请参见:这个答案

当你构造一个tf时,你可以设置GPU内存的分配比例。会话通过传递一个tf。GPUOptions作为可选配置参数的一部分:

# Assume that you have 12GB of GPU memory and want to allocate ~4GB:
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)

sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

per_process_gpu_memory_fraction充当同一台机器上每个GPU上的进程将使用的GPU内存量的硬上限。目前,这个分数统一应用于同一台机器上的所有gpu;没有办法在每个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.0和2.1版本,请使用以下代码片段:

 import tensorflow as tf
 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
 tf.config.experimental.set_memory_growth(gpu_devices[0], True)

对于以前的版本,下面的代码段用于我:

import tensorflow as tf
tf_config=tf.ConfigProto()
tf_config.gpu_options.allow_growth=True
sess = tf.Session(config=tf_config)

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

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