Блог

Я реализовал распределенную стратегию для обучения моей модели на нескольких графических процессорах.

 strategy = tf.distribute.MirroredStrategy(devices=devices[:FLAGS.n_gpus])
strategy.run(fn=self.train_step, args=(model, data))
 

Моя модель теперь стала более сложной и крупной, и мне пришлось уменьшить размер пакета, чтобы он поместился на графических процессорах.
Градиент сейчас довольно шумный, и я хочу снова увеличить размер пакета, накапливая градиенты.

Теперь мой вопрос: возможно ли это вообще при использовании зеркальной стратегии? Я знаю, что потери и градиенты в любом случае объединяются по репликам, так есть ли способ суммировать их по репликам И, например, цикл, выполняемый по пакетам? Я попробовал простую вещь и вернул вычисленные градиенты для каждой реплики, чтобы добавлять и применять их за пределами strategy.run() подобного:

 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)
 

но Tensorflow говорит мне, что это невозможно, и градиенты должны применяться внутри strategy.run() . Это также имеет смысл для меня, но мне интересно, есть ли возможность накапливать градиенты И использовать зеркальную стратегию?

Вы могли бы использовать tf.distribute.ReplicaContext.all_reduce : Это отличается от Strategy.reduce того, что оно предназначено для контекста реплики и не копирует результаты на хост-устройство. all_reduce обычно следует использовать для сокращений внутри шага обучения, таких как градиенты.

Более подробную информацию можно найти в документе здесь.