python - 如何引发 tensorflow 内存不足错误的异常？

Question

我正在循环中使用 sess.run() 运行多个 tensorflow 推理，并且某些推理对于我的 GPU 来说过于繁重。

我收到如下错误:

2019-05-23 15:37:49.582272: E tensorflow/core/common_runtime/executor.cc:623] 
Executor failed to create kernel. Resource exhausted: OOM when allocating tensor of shape [306] and type float

我希望能够捕获这些特定的 OutOfMemory 错误，但不能捕获其他错误(这可能是由于错误的输入格式或损坏的图表所致。)

显然，结构类似于:

try:
   sess.run(node_output, feed_dict={node_input : value_input})
except:
    do_outOfMemory_specific_stuff()

不起作用，因为其他类型的错误将导致调用 do_outOfMemory_specific_stuff 函数。

知道如何捕获这些 OutOfMemory 错误吗？

Answer 1

您应该能够通过以下方式捕获它:

...
except tf.errors.ResourceExhaustedError as e:
    ...

根据this documentation .