Как выбрать / настроить AWS экземпляров графического процессора для ускорения работы TensorFlow.keras?

Question

У меня есть модель tf.keras LSTM с примерно 600 МБ обучающих данных. На каждую тренировочную эпоху уходит около 90 секунд. У меня последняя версия tensorflow - v2.2. Он работает на экземпляре AWS g3.4xlarge. Этот экземпляр оснащен графическим процессором Tesla M60 от Nvidia и 8 ГБ оперативной памяти для графического процессора.

Я хочу выполнить настройку гиперпараметров, поэтому мне нужно ускорить выполнение. Поэтому я переместил модель и данные в экземпляр AWS p3.2xlarge, который имеет графический процессор P100 с 16 ГБ оперативной памяти. Затем я обнаружил, что время обучения для каждой эпохи вообще не изменилось.

Итак, я переключился на еще больший AWS экземпляр, p3.8xlarge, который имеет 4 GPU Tesla V100 и 64 ГБ оперативной памяти. При первом запуске он использовал только 1 графический процессор и давал одинаковое время выполнения для каждой эпохи, около 90 секунд. Итак, я нашел статью на веб-сайте tenorflow о том, как использовать все графические процессоры, https://www.tensorflow.org/guide/gpu

Таким образом, при работе всех 4 графических процессоров время выполнения для эпохи увеличилось с 90 секунд до 112 секунд. ! Я использовал nvtop и nvidia-smi для мониторинга производительности графического процессора, как показано ниже.

Что мне нужно сделать, чтобы сократить время выполнения?

Answer 1

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

Прежде всего, если вы можете изменить свою сеть, чтобы уменьшить количество параметров, это будет здорово. Итак, уменьшение вашей модели, очевидно, является первым, что заставит ее работать быстрее.

Но, если это невозможно, вот две вещи, которые вы можете сделать:

1. Используйте tf. mixed_precision для ускорения работы модели.

from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

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

Используйте XLA.

import tensorflow as tf

tf.config.optimizer.set_jit(True)

Accelerated Linear Algebra, XLA - это компилятор c, специфичный для домена, для матричных операций. Это может сделать вашу сеть быстрее без каких-либо изменений в исходном коде.

Попробуйте оба из них. Я лично использовал смешанную точность, и это, безусловно, сокращает время сети.

Также, пожалуйста, в следующий раз не переходите на более крупные инстансы, так как это пустая трата ваших денег. Просто попробуйте уменьшить количество параметров (т.е. размер сети) или воспользуйтесь этими двумя приемами. Я обновлю этот ответ, если найду новый трюк.