Авторы: Тоби Бойд, Янан Цао, Санджой Дас, Томас Йорг, Джастин Лебар
XLA - это компилятор графиков TensorFlow, который вы можете использовать для ускорения ваших моделей TensorFlow ML уже сегодня с минимальными изменениями исходного кода. В этом посте описывается, что такое XLA, и показано, как вы можете опробовать его на собственном коде.
TensorFlow 1.12 (с XLA) обеспечивает значительный прирост производительности по сравнению с TF 1.11 (без XLA) при обучении ResNet50 v1.0 на графических процессорах NVIDIA® Tesla® V100: 10526 изображений в секунду с синтетическими данными и 10267 изображений в секунду с реальными данными (см. Приложение для инструкции по воспроизведению). Мы наблюдали ускорение в диапазоне от 1,13x до 3,04x на различных внутренних моделях.
XLA: TensorFlow, скомпилировано!
Обычно, когда вы запускаете граф TensorFlow, все операции выполняются индивидуально исполнителем графа TensorFlow. Каждая операция имеет предварительно скомпилированную реализацию ядра графического процессора (поставляемую как часть двоичного файла TensorFlow), которой отправляет исполнитель графа.
XLA предоставляет альтернативный режим запуска моделей TF: он компилирует ваш график TensorFlow в последовательность ядер графического процессора, сгенерированных специально для вашей модели. Поскольку эти ядра уникальны для вашей программы, они могут использовать информацию, относящуюся к конкретной модели, для оптимизации.
В качестве примера рассмотрим оптимизацию, которую выполняет XLA в контексте простого вычисления TensorFlow:
При работе без XLA график запускает три ядра: одно для умножения, одно для сложения и одно для сокращения.
Однако XLA может оптимизировать граф так, чтобы он вычислял результат за один запуск ядра. Это достигается за счет «объединения» сложения, умножения и сокращения в одно ядро графического процессора. Более того, эта объединенная операция не записывает в память промежуточные значения, полученные с помощью y * z и x + y * z; вместо этого он «передает» результаты этих промежуточных вычислений непосредственно пользователям, сохраняя их полностью в регистрах графического процессора.
Fusion - это самая важная оптимизация XLA. Пропускная способность памяти обычно является самым ограниченным ресурсом аппаратных ускорителей, поэтому удаление операций с памятью - один из лучших способов повысить производительность.
Использование XLA в ваших моделях
XLA предоставляет API, xla.compile, который позволяет вам явно вызывать компилятор XLA в части вашего графа TensorFlow. xla.compile принимает функцию Python, которая генерирует вычисление TensorFlow и связывает сгенерированные вычисления для компиляции с помощью XLA. xla.compile возвращает список тензоров, каждый из которых соответствует результату вычисления, созданного переданной функцией, но теперь оптимизированного XLA.
Таким образом, вычисление, сгенерированное приведенной выше моделью model_fn, может быть выполнено с помощью XLA, вызвав xla.compile следующим образом:
Вы можете использовать флаг командной строки (или другую произвольную логику), чтобы контролировать, компилируются ли ваши вычисления с помощью XLA или нет. Модели xla.compile часто называют
что позволяет легко экспериментировать.
Мы создали колаб, в котором вы можете поиграть с xla.compile на чуть более сложной модели.
xla.compile - не единственный способ вызвать XLA для подграфа TensorFlow; в частности, есть способы попросить TensorFlow автоматически найти XLA-совместимые подграфы и скомпилировать их с помощью XLA, но мы не будем обсуждать их в этой публикации.
Предостережения при использовании XLA
Во-первых, серверная часть графического процессора XLA в настоящее время является экспериментальной - хотя мы не знаем о каких-либо серьезных проблемах, она не тестировалась при широком производственном использовании.
Во-вторых, xla.compile еще не работает с высокоуровневыми API Keras, такими как model.fit (хотя вы можете использовать Keras ops) или в режиме ожидания. Мы активно работаем над API, чтобы включить XLA в этих режимах; Будьте на связи.
В-третьих, XLA не может компилировать все графики TensorFlow; в xla.compile можно передавать только графики со следующими свойствами.
Все операции должны иметь выводимые формы
XLA должна иметь возможность выводить формы для всех операций, которые он компилирует, с учетом входных данных для вычислений. Таким образом, модельная функция, которая создает тензор непредсказуемой формы, завершится ошибкой при запуске. (В этом примере форма вывода tf.expand_dims зависит от random_dim_size, который не может быть выведен с учетом x, y и z.)
Обратите внимание, что, поскольку XLA является JIT-компилятором, формы могут различаться в разных прогонах, если они могут быть выведены с учетом входных данных кластера. Так что этот пример подходит.
Все операции должны поддерживаться XLA.
Не все операции TensorFlow могут быть скомпилированы с помощью XLA, и если в вашей модели есть операция, которую XLA не поддерживает, компиляция XLA завершится ошибкой. Например, XLA не поддерживает операцию tf.where, поэтому, если ваша модельная функция включает эту операцию, она завершится ошибкой при запуске с xla.compile.
Каждая операция TensorFlow, поддерживаемая XLA, имеет вызов REGISTER_XLA_OP в tensorflow / compiler / tf2xla / kernels /, поэтому вы можете использовать grep для экземпляров макроса REGISTER_XLA_OP, чтобы найти список поддерживаемых операций TensorFlow.
Приложение
Производительность по тестам Google
Ниже приведен график относительного ускорения / замедления TensorFlow с XLA и TensorFlow без XLA на всех тестовых моделях команды XLA, работающих на графическом процессоре V100. Мы ничего не сдерживаем; это полный набор тестов, которые мы используем сегодня при оценке компилятора.
Каждая полоса представляет собой полную модель, например «Resnet50 обучающих изображений / сек» или «скорость вывода по внутренней модели Google». Ось X отсортирована по ускорению.
Ваш опыт может отличаться, особенно с учетом того, что мы оптимизировали XLA специально на основе многих из этих тестов! Тем не менее многие из них хорошо себя зарекомендовали, и мы продолжаем совершенствоваться.
Воспроизведение теста ResNet50 v1.0
В разделах ниже описывается настройка экземпляра Google Cloud и выполнение теста ResNet50.
Подготовьте данные
Этот шаг необходим только для проверки реальных данных и может занять несколько часов. Мы рекомендуем делать это на инстансе только с ЦП, чтобы снизить затраты на вычисления. Используя инструкции для imagenet_to_gcs.py, создайте данные imagenet в формате TFRecord и поместите их в Google Cloud Storage Bucket.
Создать экземпляр GCE
Приведенный ниже фрагмент создает экземпляр виртуальной машины Google Deep Learning на облачной платформе Google с восемью графическими процессорами Tesla® V100.
