Python-如何防止张量流分配GPU内存的总量？

Question

我在共享计算资源的环境中工作，即，我们有几台服务器计算机，每台服务器计算机都配备了一些Nvidia Titan X GPU。

对于中小型模型，Titan X的12 GB通常足以让2-3个人在同一GPU上同时进行训练。如果模型足够小，以至于单个模型无法充分利用GPU的所有计算单元，那么与运行另一个训练过程相比，这实际上可以提高速度。即使在同时访问GPU确实减慢了单个训练时间的情况下，具有让多个用户同时在GPU上进行训练的灵活性仍然很好。

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

如果有人知道对于一个给定的模型就足够了，有没有办法使TensorFlow只分配4 GB的GPU内存？

Answer 1

你可以tf.Session通过tf.GPUOptions在可选config参数中传递a 来设置构造a时要分配的GPU内存的比例：

# 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进行设置。