¿Cómo reduzco el consumo de memoria en un bucle en TensorFlow?

Tengo un bucle en TensorFlow que se ve así:

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) 

Recibo un error de OOM al asignar los gradientes para esta capa, ya que cada multiplicación de matrices es una operación diferente en el gráfico que toma memoria. ¿Hay alguna forma de evitar que TensorFlow asigne todas estas operaciones al mismo tiempo?

Este es un gráfico desafiante para que TensorFlow optimice, ya que las activaciones de cada capa deben mantenerse para agregar un gradiente único para W Una posibilidad es pasar el argumento de aggregation_method experimental al llamar a optimizer.optimize() .

Por ejemplo, puedes probar lo siguiente:

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

Esta opción agrega con entusiasmo los gradientes para las variables de uso recurrente en su lugar, en lugar de mantenerlos todos en la memoria hasta que se hayan computado todos los gradientes. Si esto no funciona, el tf.AggregationMethod.EXPERIMENTAL_TREE puede funcionar mejor.