Попытка ошибки памяти Tensorflow GPU - за исключением того, что ошибка не перехватывается

Question

Я пытаюсь запустить оптимизацию гиперпараметра (используя мяту) в большой сети с большим количеством обучаемых переменных. Я обеспокоен тем, что, когда я пробую сеть с слишком большим количеством скрытых модулей, Tensorflow выдаст ошибку памяти GPU.

Мне было интересно, есть ли способ отловить ошибку памяти графического процессора, выдаваемую Tensorflow, и пропустить пакет гиперпараметров, который вызывает ошибку памяти.

Например, я хотел бы что-то вроде

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

Когда я пытаюсь протестировать исключение ошибки памяти выше, код ломается и просто печатает ошибку памяти графического процессора и не переходит в блок исключений.

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/Hak333m/stuff/blob/master/gpu_oom.py