Come ridurre il consumo di memoria in un ciclo in TensorFlow?

Question

Ho un ciclo in tensorflow che assomiglia a questo:

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) 

sto ottenendo un errore quando OOM ripartizione dei gradienti di questo strato, poiché ogni moltiplicazione matrice è un'operazione differente nel prendere memoria grafico. C'è un modo per impedire a TensorFlow di allocare tutte queste operazioni contemporaneamente?

Answer 1

Questo è un grafico stimolante per ottimizzare TensorFlow, poiché le attivazioni da ogni livello devono essere mantenute per aggregare un singolo gradiente per W. Una possibilità è passare l'argomento experimentalaggregation_method quando si chiama optimizer.optimize().

Ad esempio, è possibile provare i seguenti:

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

Questa opzione aggrega avidamente i gradienti per le variabili ricorrentemente utilizzati sul posto, piuttosto che tutti mantenere in memoria fino a quando tutte le pendenze sono stati calcolati. Se questo non funziona, lo tf.AggregationMethod.EXPERIMENTAL_TREE potrebbe funzionare meglio.