如何引发张量流内存不足错误的异常

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:
    ...

根据此文档。