Acumula gradientes con estrategia distribuida en Tensorflow 2

15 Jul, 2023 Programación 0

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

July 15, 2023 at 1:12 am

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 llamada strategy.run(). Aquí tienes un ejemplo de cómo puedes acumular gradientes usando una estrategia espejada:

strategy = tf.distribute.MirroredStrategy(devices=devices[:FLAGS.n_gpus])

# Crea una lista vacía para almacenar los gradientes por lote
gradient_accumulators = [tf.Variable(tf.zeros_like(tv.initialized_value()), trainable=False) for tv in model.trainable_variables]

# Define la función del paso de entrenamiento
@tf.function
def train_step(inputs):
    images, labels = inputs

    # Define la función para calcular los gradientes por un solo lote
    def step_fn(inputs):
        images, labels = inputs
        with tf.GradientTape() as tape:
            logits = model(images)
            loss = compute_loss(labels, logits)
        gradients = tape.gradient(loss, model.trainable_variables)
        return gradients

    # Ejecuta la función de paso en todas las réplicas
    per_replica_gradients = strategy.experimental_run_v2(step_fn, args=(inputs,))

    # Acumula los gradientes en todos los lotes
    for i, gradients in enumerate(per_replica_gradients):
        for j, gradient in enumerate(gradients):
            gradient_accumulators[j].assign_add(gradient)

    # Solo aplica los gradientes cada n lotes
    if tf.equal(tf.math.mod(optimizer.iterations, gradient_accumulation_steps), 0):
        gradients = [gradient_accumulator / gradient_accumulation_steps for gradient_accumulator in gradient_accumulators]
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        # Reinicia los acumuladores
        for gradient_accumulator in gradient_accumulators:
            gradient_accumulator.assign(tf.zeros_like(gradient_accumulator))

# Define el flujo de entrada y el conjunto de datos
train_dataset = ... # Crea el flujo de entrada
train_dataset = strategy.experimental_distribute_dataset(train_dataset)

# Entrena el modelo
for epoch in range(num_epochs):
    for x in train_dataset:
        strategy.run(train_step, args=(x,))

En este ejemplo, definimos una lista gradient_accumulators que almacena un acumulador por cada variable entrenable en el modelo. En la función train_step(), definimos una función step_fn() que calcula los gradientes para un solo lote. Ejecutamos esta función en todas las réplicas utilizando strategy.experimental_run_v2(). Luego acumulamos los gradientes en todos los lotes sumándolos al acumulador correspondiente en gradient_accumulators. Solo aplicamos los gradientes cada gradient_accumulation_steps lotes, y dividimos cada acumulador por gradient_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.

es.davy.ai

¿Tienes una pregunta?

Acumula gradientes con estrategia distribuida en Tensorflow 2

Answer

davy.ai