一尘不染

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

Python

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

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

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

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


阅读 355

收藏
2020-02-14

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

2020-02-14