使用Tensorflow进行CPU / GPU内存使用

Question

我想运行一个也在服务器上使用Tensorflow的Python脚本。当我在没有会话配置的情况下运行它时，该进程分配了所有GPU内存，阻止任何其他进程访问GPU。

服务器规格如下：

• CPU：2x 12cores@2.5 GHz，
• RAM：256GB，
• 磁盘：2x 240GB SSD，6x 4TB @ 7200RPM，
• GPU：2x Nvidia Titan X。

此服务器由其他同事共享，因此我不能真正分配所有GPU内存。

在Tensorflow的网站上，我发现了这些指令，为使用过的GPU内存设置了一个阈值。

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

我有两个问题： 1.如果分配的GPU内存不足，那么进程会自动使用CPU，还是会崩溃？ 2.如果进程想要使用GPU但GPU已经完全分配会发生什么？

谢谢。

Answer 1

1. 如果分配的GPU内存不足，TF将会出现Out Of Memory错误并崩溃。

2. 在这种情况下，TF也会崩溃。

Answer 2

Tensorflow提供了一些选项作为其分配所有可用GPU内存的默认行为的替代选择（它可以避免内存碎片并提高运行效率）。这些选项是：

• config.gpu_options.allow_growth - 配置为True时会根据需要动态分配更多内存，但永远不会释放内存
• config.gpu_options.per_process_gpu_memory_fraction - 当配置为0到1之间的双精度时，将静态分配可用内存的一部分而不是所有内存

有关详细信息，请参阅https://www.tensorflow.org/tutorials/using_gpu#allowing_gpu_memory_growth。