TensorFlow中的RNN模型内存不足

我使用TensorFlow中的rnn.rnn助手实现了一个序列到序列模型

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)

在为LSTM单元分配梯度时，该模型在具有16 GB内存的Titan X上内存不足：

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正试图同时为所有时间步分配梯度。有没有避免这种情况的方法？

函数

tf。梯度

以及优化器的

minimize

方法允许您设置名为

aggregation\u method

的参数。默认值为

ADD\N

。该方法构造图形时，需要同时计算所有梯度

还有另外两种未记录的方法称为

tf.AggregationMethod.EXPERIMENTAL\u TREE

和

tf.AggregationMethod.EXPERIMENTAL\u acculate\u N

，它们没有此要求