python - Tensorflowcompute_gradients 和 apply_gradients 内存不足

Question

我有以下几行作为程序的一部分:

tensor_gradients = optimizer.compute_gradients(cross_entropy)

with tf.Session() as session:
    for step in range(20000):
        batch = mnist.train.next_batch(train_batch_size)
        feed = {input_x: batch[0], input_y: batch[1]}
        gradients = session.run([tensor_gradients], feed)[0]
        for i in range(len(gradients)):
            gradients[i] = (gradients[i][0], tensor_gradients[i][1])
        ... computation on gradients ...
        training_step = optimizer.apply_gradients(gradients)
        training = session.run([training_step], feed)

我这样做的原因是因为我想使用 numpy 修改渐变。上面的代码在第 800 步左右耗尽了内存。但是，如果将 optimizer.apply_gradients 步骤替换为 tensor_gradients，则代码不会耗尽内存。

        training_step = optimizer.apply_gradients(tensor_gradients)

对可能发生的事情有什么想法吗？除了上面的行之外，其余代码保持不变。梯度中的 numpy 数组是否可能没有被垃圾收集，因为它们被传递到 apply_gradients 步骤中？我不知道内存泄漏可能在哪里，或者我是否通过将修改后的梯度(以 numpy 数组形式)传递回 apply_gradients 无意中添加到 tensorflow 图。

Answer 1

Any ideas at what might be happening?

发生 OOM 是因为您在循环内构建图表:这会构建一个包含 20,000 个节点的图表，并且运行它可能需要比您拥有的内存更多的内存。

将构建图形的所有 TF 操作移到循环之外，即除了 feed_dict 构造和 sess.run 调用之外的所有操作。

<小时/>

回复评论

Apply gradients builds the graph?

是的，如果你查看文档:

Returns:
  An `Operation` that applies the specified gradients. If `global_step`
  was not None, that operation also increments `global_step`.