使用急切执行时,如何防止Tensorflow分配GPU内存的总量?

时间:2018-11-02 19:27:26

标签: tensorflow

我有一个已经回答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 个答案:

答案 0 :(得分:1)

您使用的是哪个tensorflow版本? 我使用的是在Ubuntu 16.04上从源代码(r1.12分支)构建的tensorflow v1.12,具有Cuda 9.2和Cudnn 7.2.1,并且allow_growth选项似乎可以按期望的方式正常工作。