Модель RNN, заканчивающаяся памятью в TensorFlow

Question

Я реализовал модель Sequence to Sequence, используя хинтер rnn.rnn в TensorFlow.

with tf.variable_scope("rnn") as scope, tf.device("/gpu:0"): 
    cell = tf.nn.rnn_cell.BasicLSTMCell(4096) 
    lstm = tf.nn.rnn_cell.MultiRNNCell([cell] * 2) 

    _, cell = rnn.rnn(lstm, input_vectors, dtype=tf.float32) 
    tf.get_variable_scope().reuse_variables() 
    lstm_outputs, _ = rnn.rnn(lstm, output_vectors, initial_state=cell) 

Модель работает из памяти на Titan X с 16 Гб памяти при выделении градиентов для LSTM клеток:

W tensorflow/core/kernels/matmul_op.cc:158] Resource exhausted: OOM when allocating tensor with shape[8192,16384] 
W tensorflow/core/common_runtime/executor.cc:1102] 0x2b42f00 Compute status: Resource exhausted: OOM when allocating tensor with shape[8192,16384] 
    [[Node: gradients/rnn/RNN/MultiRNNCell_1/Cell0/BasicLSTMCell/Linear/MatMul_grad/MatMul_1 = MatMul[T=DT_FLOAT, transpose_a=true, transpose_b=false, _device="/job:localhost/replica:0/task:0/gpu:0"](rnn/RNN/MultiRNNCell_1/Cell0/BasicLSTMCell/Linear/concat, gradients/rnn/RNN/MultiRNNCell_1/Cell0/BasicLSTMCell/add_grad/tuple/control_dependency)]] 

Если я уменьшить длину входных и выходных последовательностей 4 или менее модель работает без проблем.

Это указывает на то, что TF пытается выделить градиенты для всех временных шагов одновременно. Есть ли способ избежать этого?

Answer 1

Функция tf.gradients, а также метод оптимизаторов minimize позволяют задать параметр aggregation_method. Значение по умолчанию: ADD_N. Этот метод строит график таким образом, что все градиенты должны вычисляться одновременно.

Существует два других недокументированных метода: tf.AggregationMethod.EXPERIMENTAL_TREE и tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N, которые не имеют этого требования.