如何在TensorFlow循环中减少内存消耗？

Question

我在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同时分配所有这些操作？

Answer 1

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

例如，您可以尝试以下方法：

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可能会运行得更好。