我使用 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.gradients以及优化器的minimize方法允许您设置名为aggregation_method的参数。默认值为ADD_N。该方法以需要同时计算所有梯度的方式构建图。
还有另外两个未记录的方法，称为 tf.AggregationMethod.EXPERIMENTAL_TREE 和 tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N，它们没有此要求。
关于tensorflow - TensorFlow 中的 RNN 模型内存不足，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/36139889/