如何减少 TensorFlow 循环中的内存消耗？答案

【问题标题】：How I reduce memory consumption in a loop in TensorFlow?如何减少 TensorFlow 循环中的内存消耗？
【发布时间】：2016-07-11 16:50:17
【问题描述】：

我在 TensorFlow 中有一个如下所示的循环：

with tf.device("/gpu:1"):
    losses = []

    for target, output in zip(targets, lstm_outputs):
        logits = tf.matmul(W, output) + b
        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits, target)
        losses.append(loss)

    total_loss = tf.add_n(losses)

在为该层分配梯度时，我遇到了 OOM 错误，因为每个矩阵乘法都是图形中占用内存的不同操作。有没有办法阻止 TensorFlow 同时分配所有这些操作？

【问题讨论】：

标签： python gpu tensorflow

【解决方案1】：

对于 TensorFlow 而言，这是一个难以优化的图，因为必须保留来自每一层的激活来聚合 W 的单个梯度。一种可能性是在调用optimizer.optimize() 时传递experimental aggregation_method 参数。

例如，您可以尝试以下操作：

optimizer = tf.train.AdagradOptimizer(...)  # Or another optimization algorithm.
train_op = optimizer.minimize(
    total_loss,
    aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)

此选项会急切地聚合经常使用的变量的梯度，而不是在计算完所有梯度之前将它们全部保存在内存中。如果这不起作用，tf.AggregationMethod.EXPERIMENTAL_TREE 可能会更好。

【讨论】：

我已经尝试了这两个 EXPERIMENTAL_ACCUMULATE_N 和 EXPERIMENTAL_TREE 无济于事。我会尝试使用 while 循环。
我能够通过从稳定版本更新到主版本并结合使用 EXPERIMENTAL_ACCUMULATE_N 来解决这个问题。 @mrry 感谢您的努力和响应。