Tensorflow GPU内存报错try-except not catching the error

Question

我正在尝试运行在具有大量可训练变量的大型网络上进行超参数优化（使用留兰香）。我担心当我尝试隐藏单元数量过多的网络时，Tensorflow 会抛出 GPU 内存错误。

我想知道是否有一种方法可以捕获 Tensorflow 抛出的 GPU 内存错误并跳过导致内存错误的那批超参数。

例如，我想要这样的东西

import tensorflow as tf 

dim = [100000,100000]
X   = tf.Variable( tf.truncated_normal( dim, stddev=0.1 ) )

with tf.Session() as sess:
    try:
        tf.global_variables_initializer().run()
    except Exception as e :
        print e

当我在上面尝试测试内存错误异常时，代码中断并仅打印 GPU 内存错误并且没有前进到 except 块。

Answer 1

试试这个：

import tensorflow as tf

try:
    with tf.device("gpu:0"):
        a = tf.Variable(tf.ones((10000, 10000)))
        sess = tf.Session()
        sess.run(tf.initialize_all_variables())
except:
    print("Caught error")
    import pdb; pdb.set_trace()

来源：https://github.com/yaroslavvb/stuff/blob/master/gpu_oom.py

Tensorflow GPU内存报错try-except not catching the error

Tensorflow GPU memory error try-except not catching the error

python

gpu

except

deep-learning

tensorflow