TensorFlow强制使用GPU

Question

我们的GPU处于独占模式。有时，某些用户可能会手动登录计算机并窃取GPU。

如果TensorFlow脚本中的GPU初始化失败，如何引发异常？我注意到当TensorFlow无法初始化GPU时，它会输出错误消息但在CPU上运行。我想停止它而不是在CPU上运行。

Answer 1

如果您使用以下方法强制图表的任何部分在GPU上运行：

with tf.device('/device:GPU:0'):

然后你的会话变量初始化程序将停止运行并在没有GPU可用时抛出InvalidArgumentError。