ПОСТЫ

Ускоряем TensorFlow сборкой из исходников

February 26, 2020

Как заставить TensorFlow работать на CPU до 7 раз быстрее.

Или примерно в 3 в моем случае.

Мой кейс

Должен заметить что это первая статья связанная с ML(надеюсь будет больше)

Так уж вышло что в данный момент я работаю над проектом, частью которого является небольшая нейронная сеть. А точнее две: по-больше по-важнее и по-меньше по-необязательнее. Речь будет идти в основном о первой.

В общем-то это известная в узких кругах модель: YOLO/darknet, а точнее вторая ее версия. Казалось бы при чем тут Tensorflow? А при том, что использую я ее не напрямую(потому что ~~пошли они~~ неочевидно как пользоваться при отсутствующей документации), а через форк darkflow, который “дословно” переводит darknet в tensorflow понятные штуки.

Статка

Еще до запуска на лептопе(да не стоило) было понятно, что нужно делать сервак, однако я решил попробовать. Итак результат на ноутбуке(0)

№	Dataset size	Batch	t₁₀₀₀	t₁	CPU Load	t	CPU	RAM	OS
0	100	16	5ч	18с	100%	105°C	2 x 2.7GHz	8Gb DDR3	macOS

t₁₀₀₀ - Время на 1000 шагов обучения

t₁ - Время на 1 шаг обучения

Да, ерунда конечно, но эксперимент есть эксперимент, пришлось потерпеть вечерок шум ноутбука.

Пробуем сервер(1), результаты:

№	Dataset size	Batch	t₁₀₀₀	t₁	CPU Load	t	CPU	RAM	OS
0	100	16	5ч	18с	100%	105°C	2 x 2.7GHz	8Gb DDR3	macOS
1	100	16	3ч	11с	100%	n/a	8 x 2.3GHz	16Gb DDR3	Ubuntu

Прирост конечно есть, но все еще до невозможности медленно. Да еще и куча варнингов, мол можно быстрее если скомпилировать с каким-то AVX и еще кучей библиотек. Штош, пробуем(2):

№	Dataset size	Batch	t₁₀₀₀	t₁	CPU Load	t	CPU	RAM	OS
0	100	16	5ч	18с	100%	105°C	2 x 2.7GHz	8Gb DDR3	macOS
1	100	16	3ч	11с	100%	n/a	8 x 2.3GHz	16Gb DDR3	Ubuntu
2	100	16	40м	3с	80%	n/a	8 x 2.3GHz	16Gb DDR3	Ubuntu

Что? Прирост производительности в 4 раза? Да еще и виртуалка перестала сыпать предупреждениями об использовании проца на 100%. Магия не иначе. Разбираемся.

Почему так

Давай взглянем что я там накомпилировал. Я использовал следующие флаги при сборки Bazelisk-ом:

--copt=-mavx --copt=-mavx2 --copt=-mfma

За каждым из них скрывается библиотека. Посмотрим поближе что они делают:

AVX и AVX2

AVX это аббревиатура от Advanced Vector Extensions. Это расширение команд для процессоров x86, которое вкючает в себя кучу различных обновлений в схеме работы с памятью в целом и во время вычислений, определяет новые регистры(или тип того), но самое главное делает работу с векторной(тензорной) ~~математикой~~ арифметикой быстрее и проще. А как мы знаем, TensorFlow, неожиданно, работает именно с ней.

mFMA

Fused Multiply-Add - так же набор инструкций для процессоров x86, упрощающий и ускоряющий работу со сложением-умножением чисел с плавающей запятой.

Что-нибудь еще и GPU сюда же

Чтобы понять что ставить, сделай cat /proc/cpuinfo | grep flags и сравни это с тем, что поддерживает Bazel сборка TensorFlow. Или еще лучше просто поставь собранный TensorFlow каким-нибудь pip install tensorflow и выполни python -c 'import tensorflow as tf; tf.Session()' и tf сам скажет тебе какие инструкции поддерживает твой процессор, но библиотеки не оптимизированы под них.

Есть легенда, что если поставить TensorFlow через Anaconda(conda), то можно тоже добится неплохого прироста без того, что будет описано дальше в статье, однако лично мне это ощутимого прироста не дало, а варнинги остались на своем месте

Если у тебя есть GPU - на ноуте или тем более сервере(если так, то не знаю что ты тут забыл, но раз уже забыл, то поделись🙏🏻).

Эта установка заставит тебя попотеть

Лично у меня, а я себя тупым не считаю, получилось примерно с 5 раза. Почему сейчас объясню.

Ставим ненавистный Bazelisk(Bazel)

Ахтунг: бомбеж

Я не знаю кто придумал это, но для него есть отдельный котел в аду, ведь человек, или люди, который решил использовать вместо make/cmake - это, просто-напросто садист. Начнем с того, что в зависимости от того, какую версию TF ты ставишь, тебе нужны не просто разные версии TF, а разные версии установщика, а иногда еще и JDK. Продолжим это тем, что сборка конфигурируется и скриптом и флагами, и где что и как друг на друга влияет вообще не ясно. Установщик этого установщика существует минимум в трех вариантах и ни один их них не рабочий. В общем вбив в гуглий после 4-ой, неудачной попытки What the fuck is Bazel?!, я не был удивлен статьями вроде этой:

За дело

Итак узнаем какую версию bazel нам нужно поставить для нужной нам версии tf и поставим её.

Для этого открывай репозиторий tensorflow и переходи в релиз-ветку по тегу.
В файле configure.py ищем _TF_MIN_BAZEL_VERSION и/или _TF_MAX_BAZEL_VERSION
ЕСЛИ такого файла нет(что обычно бывает в версиях как минимум до 2.0), то идешь сюда, ищешь свою платформу, версию tf и соответствующую версию bazel.
Ищешь нужную версию bazel в релизах на github
Качаешь wget-ом bazel-<чето там>-installer.sh, а если тоже на Ubuntu, то конечно .deb пакет
Запускаешь инсталлер один раз так, другой раз с ключем –user, либо просто apt install ./bazel-<smth>.deb
Проверь установку bazel version - он прогрузится и доставит нужные и ненужные зависимости и покажет версию. Убедись что она совпадает с той, которая тебе нужна

???
PROFIT

Готовимся к сборке

TF ты уже склонировал, и коли не дурак то сделал checkout на твою ветку. Если нет, то сделай это.

cd tensorflow
./configure - Если охота и возможность можно даже проследовать нестандартным путем и наставить всякого, что может еще ускорить tf в дальнейшем. А может и нет. А может и да.

Собираем

Как я уже писал, я собирал с AVX, AVX2 и mFMA, поэтому для меня команда выглядела так:

bazel build --config=opt --copt=-mavx --copt=-mavx2 --copt=-mfma --copt=-mfpmath=both -k --verbose_failures //tensorflow/tools/pip_package:build_pip_package

Пару часов/суток и сборка завершится. Останется только установить все это дело в Python Для этого выполни:

bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

И установи получившийся PIP пакет:

pip install /tmp/tensorflow_pkg/tensorflow-version-tags.whl

Готово.

Важное замечание

Bazel помимо остальных жирных минусов имеет еще один жирнющий: Он жрет памяти больше чем 200 хромов. Так один раз он сожрал все данные ему 16Гб оперативной памяти и благополучно повесил сервер в OOM. Так что оперативку брать надо с запасом(по совету проверенных комрадов из Google минимум 16), и если версия bazel достаточно свежая, то использовать ключ --local_ram_resources=4096 подставьте свое доступное количество памяти, тогда bazel будет пытаться себя ограничивать в данных рамках. Однако выходит у него не всегда ._.

Небольшое пояснение

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

Спасибо за внимание!