如何防止张量流分配GPU内存的总量?

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

当前回答

当你构造一个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基础上设置这个。

2015-12-10 11:00:19

其他回答

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

2019-10-05 21:05:50

我尝试在voc数据集上训练unet，但由于图像大小巨大，内存结束。我尝试了上面所有的技巧，甚至尝试了batch size==1，但没有任何改善。有时候TensorFlow版本也会导致内存问题。尝试使用

PIP install tensorflow-gpu==1.8.0

2018-10-16 06:05:52

这段代码对我有用:

import tensorflow as tf
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.InteractiveSession(config=config)

2020-07-08 22:30:04

上面所有的答案都假设使用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)

2018-02-08 03:25:04

以下是《TensorFlow深度学习》一书的节选

In some cases it is desirable for the process to only allocate a subset of the available memory, or to only grow the memory usage as it is needed by the process. TensorFlow provides two configuration options on the session to control this. The first is the allow_growth option, which attempts to allocate only as much GPU memory based on runtime allocations, it starts out allocating very little memory, and as sessions get run and more GPU memory is needed, we extend the GPU memory region needed by the TensorFlow process.

1)允许增长:(更灵活)

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

第二个方法是per_process_gpu_memory_fraction选项，它决定每个可见GPU应该分配的内存总量的百分比。注意:不需要释放内存，这样做甚至会恶化内存碎片。

2)分配固定内存:

每个GPU只分配40%的内存:

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config, ...)

注意: 不过，只有当你真的想绑定TensorFlow进程上可用的GPU内存数量时，这才有用。

2018-01-11 18:57:16

如何防止张量流分配GPU内存的总量?

推荐文章

最新文章

标签