How I reduce memory consumption in a loop in TensorFlow?

Viewed 4482

I have a loop in TensorFlow that looks like this:

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)

I am getting an OOM error when allocating the gradients for this layer, since each matrix multiplication is a different operation in the graph taking memory. Is there a way of preventing TensorFlow from allocating all these operations at the same time?

1 Answers
Related