如何防止 Tensorflow 在使用 Eager Execution 时分配全部 GPU 内存？

Posted 2023-02-23

【中文标题】如何防止 Tensorflow 在使用 Eager Execution 时分配全部 GPU 内存？

【英文标题】：How to prevent Tensorflow from allocating the totality of a GPU memory when using eager execution?

【发布时间】：2019-04-07 01:30:58

【问题描述】：

我的问题与here 已经回答的问题几乎相同，但略有不同：

我正在使用带有几个 GPU 的服务器，我与同事共享这些 GPU 来训练我们的深度学习模型。服务器还应该运行一个从我们的模型中采样的小型 Web 应用程序。采样脚本使用相对较新的eager execution。从理论上讲，它允许我通过提供如下配置来阻止 Tensorflow 分配所有 GPU 内存：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
tf.enable_eager_execution(config=config)

实际上，这不起作用。急切执行的文档还指出，并非所有适用于会话的配置选项都可以在急切执行中工作 (here)。但是那我该如何限制使用的内存呢？

我知道我可以像这样限制可见设备：

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

但我不想不断地阻塞整个 GPU 来执行一个偶尔被调用并且实际上需要更少内存的任务。

【参考方案1】：

您使用的是哪个 TensorFlow 版本？ 我在带有 Cuda 9.2 和 Cudnn 7.2.1 的 Ubuntu 16.04 上使用从源代码（r1.12 分支）构建的 tensorflow v1.12，并且 allow_growth 选项似乎在急切执行时按预期工作。

【讨论】：

我做了一些实验，是的，你是对的。当我在纯 Tensorflow 中做所有事情时，它似乎有效。不过我使用 Keras，并认为它会继承该配置。现在谢谢。当我找到解决方案时，我会发布更新。