Tensorflow GPU内存在均方误差期间耗尽

时间:2017-03-27 22:35:56

标签: tensorflow gpu recurrent-neural-network

我有一个Tensorflow模型,是一个使用长期短期记忆的循环神经网络。状态大小为3000,每个输入步骤有300个输入,大约有500个时间步长,每个时间步长有1个输出。我正在训练序列到序列模型。

对于时间步长少于500的输入,它运行良好,但大约500步时,它会因以下内存不足错误而崩溃:

ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[20375,20375]
     [[Node: gradients/mean_squared_error/Mul_grad/mul_1 = Mul[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/gpu:0"](mean_squared_error/Square, gradients/mean_squared_error/Sum_grad/Tile)]]
     [[Node: gradients/MatMul_grad/tuple/control_dependency_1/_225 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/cpu:0", send_device="/job:localhost/replica:0/task:0/gpu:0", send_device_incarnation=1, tensor_name="edge_5086_gradients/MatMul_grad/tuple/control_dependency_1", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/cpu:0"]()]]

这是在拥有12GB内存的GPU上运行。

我已经尝试在我的笔记本电脑cpu上运行它,它似乎使用非常少的内存(大约1到2 GB),但它太慢了,它从来没有达到500个时间步。我正在进行一些更改,这将使其跳过500个时间步,以查看它在GPU上运行时使用了多少内存。

我的问题是:Tensorflow可能在哪里分配形状张量[20375,20375]?它似乎与tf.mean_squared_error函数有关,但这似乎不是一个需要如此大量内存的操作。

我已经尝试过减少批量大小,但是这只会将故障点推到更多的时间步骤,而且我需要多达几千个时间步,所以这似乎不是一个好的长 - 期限解决方案我更愿意找到问题的根源。

以下是均方误差的相关代码:

initial_state_tuple = tf.contrib.rnn.LSTMStateTuple(initial_state, initial_hidden_state)


# Create the actual RNN
with tf.variable_scope(VARIABLE_SCOPE, reuse=None):
    cell = tf.contrib.rnn.BasicLSTMCell(STATE_SIZE)
    rnn_outputs, finalstate = tf.nn.dynamic_rnn(cell=cell, inputs=networkinput,
                                                initial_state=initial_state_tuple)

with tf.variable_scope(VARIABLE_SCOPE, reuse=True):
    weights = tf.get_variable(name=WEIGHTS_NAME, shape=[STATE_SIZE, 1], dtype=tf.float32)
    biases = tf.get_variable(name=BIASES_NAME, shape=[1], dtype=tf.float32)

# Build the output layers
rnn_outputs_reshaped = tf.reshape(rnn_outputs, [-1, STATE_SIZE])
network_outputs = tf.sigmoid(tf.matmul(rnn_outputs_reshaped, weights) + biases)
expected_outputs_reshaped = tf.reshape(expected_outputs, [-1, 1])

# Loss mask just cancels out the inputs that are padding characters, since not all inputs have the same number of time steps
loss_mask_reshaped = tf.reshape(loss_mask, shape=[-1])

expected_outputs_reshaped = loss_mask_reshaped * expected_outputs_reshaped
network_outputs = loss_mask_reshaped * network_outputs

loss = tf.losses.mean_squared_error(labels=expected_outputs_reshaped, predictions=network_outputs)

如果您想要所有代码,可以找到here。相关函数是buildtower()和buildgraph()。在具有GPU的计算机上运行时,常量NUM_GPUS和BATCH_SIZE设置为适当的值。

更新:我替换了

loss = tf.losses.mean_squared_error(labels=expected_outputs_reshaped, predictions=network_outputs)

error_squared = tf.pow(expected_outputs_reshaped - network_outputs, 2)
loss = tf.reduce_mean(error_squared)

发生了同样的错误。我将状态大小减小到30并将批量大小减小到5,并且错误仍然发生,尽管它确实达到了大约3000个时间步长。

更新:在做了一些研究之后,我发现,在训练具有大量时间步长的RNN时,经常使用截断的反向传播。这让我相信通过大量时间步骤的反向传播本身会占用大量内存,而我的问题并不是我构建错误的图形,而是我对梯度计算的资源需求有一个基本的误解。为此,我正在努力将我的代码更改为使用截断的反向传播。我会用结果报告。

1 个答案:

答案 0 :(得分:0)

这个项目是我第一次使用机器学习和Tensorflow,经过一些研究后,似乎有一些基本的误解。

我原以为内存使用会随着数据中的时间步长而线性缩放。因为我的模型的每个其他维度(批量大小,状态大小)都很小,所以我预计在耗尽内存之前我可以完成相当多的时间步骤。然而,似乎计算梯度的内存使用量随着时间步长的数量呈指数级增长,因此无论我使用多小的状态大小和批量大小,由于时间步长很多,它最终会耗尽我的所有记忆。

为了解决这个问题,我使用的是truncated backpropagation,其中每个批次都分成了一些固定数量的时间步骤。这并不完美,因为这意味着错误只能在这么多时间步骤中传播回来。但是,基于我在网上找到的内容,它似乎运行得很好,并且没有太多其他方法可以解决内存使用问题。

正如我之前所说的,这是我第一次使用机器学习的经验,所以如果这里有任何明显的错误,请告诉我。