Python 如何在TensorFlow中减少循环中的内存消耗?

Python 如何在TensorFlow中减少循环中的内存消耗?,python,gpu,tensorflow,Python,Gpu,Tensorflow,我在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)

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

这是TensorFlow需要优化的一个具有挑战性的图表,因为必须保持每个层的激活,以便为
W
聚合一个梯度。一种可能是在调用时传递
aggregation\u方法
参数

例如,您可以尝试以下操作:

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

该选项急切地为重复使用的变量聚合梯度,而不是将它们全部保存在内存中,直到计算完所有梯度。如果这不起作用,
tf.AggregationMethod.EXPERIMENTAL\u TREE
可能会更好。

我已经尝试了这两种实验性的累积树和实验性树,但都没有用。我将尝试使用while循环。我能够解决这个问题,通过从稳定版本升级到master,并结合使用实验性的方法。@mrry感谢您的努力和响应能力。