Acumula gradientes con estrategia distribuida en Tensorflow 2
He implementado una estrategia distribuida para entrenar mi modelo en múltiples GPUs.
strategy = tf.distribute.MirroredStrategy(devices=devices[:FLAGS.n_gpus])
strategy.run(fn=self.train_step, args=(model, data))
Mi modelo ahora se ha vuelto más complejo y grande y tuve que reducir el tamaño del lote para que quepa en las GPUs.
El gradiente ahora es bastante ruidoso y quiero aumentar el tamaño del lote nuevamente acumulando gradientes.
Ahora mi pregunta es: ¿esto es incluso posible cuando se utiliza una estrategia de espejo? Sé que la pérdida y los gradientes se combinan en todas las réplicas de todos modos, ¿entonces hay alguna forma de sumarlos a través de las réplicas y, por ejemplo, en un bucle que recorre los lotes? Intenté hacer lo más sencillo y devolver los gradientes calculados por réplica para agregarlos y aplicarlos fuera de strategy.run()
de esta manera:
for b in batches:
per_replica_gradients = strategy.run(fn=self.train_step, args=(model, data))
total_gradient += per_replica_gradients
optimizer.apply_gradients(zip(total_gradient, model.trainable_variables))
Pero Tensorflow me dice que esto no es posible y que los gradientes deben aplicarse dentro de strategy.run()
. Esto también tiene sentido para mí, pero me pregunto si hay alguna posibilidad de acumular gradientes y usar una estrategia de espejo.
davy.ai
Sí, es posible acumular gradientes y usar una estrategia espejada. Una forma de hacer esto es utilizando el método
tf.distribute.Strategy.experimental_local_results()
, el cual te permite acceder a los resultados por réplica de la llamadastrategy.run()
. Aquí tienes un ejemplo de cómo puedes acumular gradientes usando una estrategia espejada:En este ejemplo, definimos una lista
gradient_accumulators
que almacena un acumulador por cada variable entrenable en el modelo. En la funcióntrain_step()
, definimos una funciónstep_fn()
que calcula los gradientes para un solo lote. Ejecutamos esta función en todas las réplicas utilizandostrategy.experimental_run_v2()
. Luego acumulamos los gradientes en todos los lotes sumándolos al acumulador correspondiente engradient_accumulators
. Solo aplicamos los gradientes cadagradient_accumulation_steps
lotes, y dividimos cada acumulador porgradient_accumulation_steps
antes de aplicar los gradientes. Luego reiniciamos los acumuladores a cero. Esto nos permite simular efectivamente tamaños de lote más grandes sin quedarnos sin memoria en una sola GPU.