TensorFlow 的 CPU/GPU 内存使用情况

Question

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

服务器规格如下：

• CPU：2x 12cores@2.5 GHz，
• 内存：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，还是会crash？ 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。