如何引发张量流内存不足错误的异常? [英] How to raise an exception for a tensorflow out of memory error?

问题描述

我正在使用 sess.run()循环运行多个张量流推论，碰巧有些推论对我的GPU来说太重了.

I am running several tensorflow inferences using sess.run() in a loop and it happens that some inferences are too heavy for my 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错误，但不能捕获其他错误(这可能是由于错误的输入格式或损坏的图形所致.)

I would like to be able to catch these specific OutOfMemory errors but not other errors (which may be due to a wrong input format or a corrupted graph.)

很显然，它的结构类似于:

Obviously, a structure similar to :

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

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

does not work since other kind of errors will lead to a call to the do_outOfMemory_specific_stuff function.

任何想法如何捕获这些OutOfMemory错误?

Any idea how to catch these OutOfMemory errors ?

推荐答案

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

You should be able to catch it via:

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

根据此文档.

这篇关于如何引发张量流内存不足错误的异常?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！