Wie reduziere ich den Speicherverbrauch in einer Schleife in TensorFlow?

Question

Ich habe eine Schleife in TensorFlow der wie folgt aussieht:

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 eine OOM Störung erhalten, wenn die Gradienten für diese Schicht Zuteilen, da jede Matrixmultiplikation eine andere Operation in dem Graph nimmt Speicher ist. Gibt es eine Möglichkeit, zu verhindern, dass TensorFlow all diese Operationen gleichzeitig zuweist?

Answer 1

Dies ist eine anspruchsvolle Grafik für TensorFlow zu optimieren, da die Aktivierungen von jeder Schicht gehalten werden müssen, um einen einzigen Gradienten für W zu aggregieren. Eine Möglichkeit besteht darin, das experimentalaggregation_method Argument beim Aufruf optimizer.optimize() zu übergeben.

Zum Beispiel könnten Sie versuchen, die folgenden:

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

Diese Option aggregiert eifrig die Gradienten für wiederkehrend verwendeten Variablen statt, anstatt sie alle in Erinnerung behalten, bis alle der Gradienten berechnet wurden. Wenn dies nicht funktioniert, funktioniert die tf.AggregationMethod.EXPERIMENTAL_TREE möglicherweise besser.