Всем привет! Меня зовут Дмитрий, я занимаюсь разработкой в области компьютерного зрения в команде MTS AI. Так исторически сложилось, что в своей работе я использую, как правило, связку устаревшей версии TensorFlow 1 и Keras. Пришло время двигаться дальше, но прежде чем полностью перейти на TensorFlow 2, я решил сравнить наиболее популярные на сегодня фреймворки глубокого обучения: TensorFlow и PyTorch. Эта статья будет полезна всем Data Scientist'ам, кто желает узнать чуть больше про особенности и производительность TensorFlow и PyTorch.

TensorFlow и Keras

TensorFlow был впервые представлен компанией Google в ноябре 2015 года. Он появился на основе собственного решения Google для задач машинного обучения. Название фреймворка отражает суть его работы: данные представляются в виде тензоров, а при выполнении операций с ними формируется граф потока данных (граф вычислений), в котором отслеживается изменение данных и поддерживается автоматическое дифференцирование. В первой версии TensorFlow граф вычислений был статическим.

Версия TensorFlow 2.x была представлена в сентябре 2019, её главным отличием от первой версии был переход от статического к динамическому графу вычислений. При этом сохранилась поддержка статического графа вычислений:

import tensorflow as tf

tf.compat.v1.disable_v2_behavior()

Keras первоначально представлял собой высокоуровневый фреймворк глубокого обучения, основанный на базе TensorFlow, Theano и CNTK. Keras был создан одним из разработчиков Google,а название является отсылкой к Одиссее. Keras был интегрирован в TensorFlow, начиная с версии 1.4.0, а в TensorFlow 2.х уже полностью интегрирован: рекомендуется использовать только tf.keras. Самостоятельная версия Keras с поддержкой нескольких фреймворков прекратила своё развитие, так как Theano и CNTK на данный момент также не развиваются.

TensorFlow предоставляет несколько полезных для разработки API: TensorFlow Lite позволяет легко оптимизировать и подготовить обученные модели для запуска на различных мобильных устройствах и устройствах IoT; TensorFlow.js позволяет использовать JavaScript для обучения моделей и запуска обученных моделей; TensorFlow Extended предоставляет множество компонентов, в том числе REST API, для реализации пайплайнов машинного обучения в готовых продуктах.

PyTorch

PyTorch был впервые представлен компанией Facebook в октябре 2016 года. PyTorch основан на Torch – Matlab-подобной библиотеки для Lua, поддержка которой практически прекратилась в 2017 году. PyTorch также поддерживает автоматическое дифференцирование и использует динамический граф вычислений.

PyTorch достаточно удобен и прост: работа с тензорами очень похожа на работу с numpy arrays, что в совокупности с динамическим графом вычислений облегчает разработку и дебагинг.

Статический и динамический графы вычислений

При использовании статического графа вычислений программа работает следующим образом. Код, описывающий различные операции с тензорами, выполняется только один раз, в результате чего формируется статический граф вычислений (объект tf.Graph() в TensorFlow 1.x). Чтобы получить результаты вычисления графа, необходимо подать данные в граф и запустить вычисления для необходимых тензоров:

# TensorFlow 1.15.0
import numpy as np
import tensorflow as tf

# Определение входного тензора в граф
inputs = tf.placeholder(tf.float32, [1, 224, 224, 3])
# Определение графа вычислений
x = tf.layers.conv2d(inputs, 64, (3, 3), strides=2, padding='same', use_bias=False)
x = tf.layers.batch_normalization(x)
output_tensor = tf.nn.relu(x)
# Использование tf.Session() для вычислений
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    outputs = sess.run([output_tensor],  feed_dict={inputs: np.random.randint(0, 256, (1, 224, 224, 3)).astype(np.float32)})

При использовании динамического графа вычисления происходят сразу при определении операций:

# TensorFlow 2.5.0
import numpy as np
import tensorflow as tf

# Определение входного тензора
inputs = np.random.randint(0, 256, (1, 224, 224, 3)).astype(np.float32)
# Определение операций и вычисление 
x = tf.keras.layers.Conv2D(64, (3, 3), strides=2, padding='same', use_bias=False)(inputs)
x = tf.keras.layers.BatchNorimalization()(x)
outputs = tf.keras.layers.Relu()(x)

Таким образом, динамический граф вычислений является более удобным:

• вычисления выполняются сразу при определении операций;

• результаты вычислений доступны на любом шаге;

• более лёгкая и гибкая разработка, упрощённый дебагинг.

Сравнение производительности TensorFlow и PyTorch

При использовании различных версий TensorFlow и PyTorch я неоднократно замечал различия в производительности, однако в гугле я практическине нашел подробной информации с числами. Из наиболее интересных данных, здесь автор вопроса сравнивал скорость обучения при использовании TensorFlow 1.х и TensorFlow 2.х., а в этой статье 2017 года представлено сравнение производительности нескольких фреймворков глубокого обучения, однако на сегодня она уже является устаревшей.

Мне, как активному пользователю этих фреймворков, наиболее интересно сравнить скорость обучения (на GPU), скорость инференса (на GPU и на CPU), а также количество выделяемой видеопамяти при обучении и инференсе.

Для сравнения я выбрал несколько классических архитектур свёрточных нейронных сетей, которые отличаются размером, количеством и типом свёрточных слоёв и функций активации. Помимо этого, я варьировал наиболее часто изменяемые параметры: размер батча (количество одновременно обрабатываемых нейронной сетью изображений) и размер входных изображений.

Поскольку конечной целью является измерение среднего времени шага обучения/инференса, данные могут быть использованы любые. Например, случайным образом сгенерированные:

import numpy as np

image_data = np.random.randint(0, 256, (steps_number, batch_size, height, width, channels))
labels_data =  np.random.randint(0, 2, (steps_number, batch_size, 1000))

1. Скорость обучения (на GPU)

Для каждой комбинации нейронной сети, размера входных изображений и размера батча было выполнено 100 шагов обучения и измерено среднее время шага обучения. Полные результаты измерений представлены в таблице 1. Прочерки в таблицах соответствуют такой комбинации параметров, при которой для процесса не хватило видеопамяти. На рисунке 1 представлены кривые, соответствующие относительному времени шага обучения фреймворков для каждой выбранной комбинации входных параметров (за единицу принято наименьшее время шага в данной комбинации).

Таблица 1. Среднее время шага обучения (в секундах).

В большинстве случаев PyTorch показывает лучшую производительность по сравнению с TensorFlow. Грубая оценка такая: скорость обучения в PyTorch 1.9.0 на 30% быстрее, чем в TensorFlow 2.5.0, и на 45% быстрее, чем в TensorFlow 1.15.0. Однако заметно, что при обучении vgg16 и efficientnet в TensorFlow 2.5.0 показывает лучшую производительность, хотя и в целом близкую к производительности в PyTorch 1.9.0. Оптимизация видеопамяти в TensorFlow 2.5.0 улучшилась по сравнению с TensorFlow 1.15.0, благодаря чему удалось запустить обучение для большего числа нейронных сетей. Оптимизация видеопамяти в PyTorch 1.9.0 выглядит чуть более превосходящей, чем в TensorFlow 2.5.0.

2. Скорость инференса (на GPU)

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

Таблица 2. Среднее время шага инференса на GPU (в секундах).

При инференсе на GPU TensorFlow 1.15.0 и PyTorch 1.9.0 показывают примерно одинаковую производительность: при грубой оценке, PyTorch 1.9.0 оказывается в среднем на 5% быстрее при использовании архитектур ResNet и VGG, а TensorFlow 1.15.0 в среднем на 5% быстрее при использовании архитектур MobileNet и EfficientNet. При этом TensorFlow 2.5.0 во всех случаях в среднем на 30% медленнее. Основная причина в этом, на мой взгляд, использование динамического графа вычислений: при переходе к статическому, скорость инференса в TensorFlow 2.5.0 становится практически такой же, как и в TensorFlow 1.15.0.

3. Скорость инференса (на CPU)

В таблице 3 представлены полные результаты измерений скорости инференса на CPU. На рисунке 3 представлены кривые, соответствующие относительному времени шагу инференса фреймворков на CPU.

Таблица 3. Среднее время шага инференса на CPU (в секундах).

Здесь результаты получились довольно неожиданные: в большинстве случаев PyTorch 1.9.0 показывает низкую производительность, причём чем сложнее нейронная сеть, тем ниже производительность. Грубая оценка скорости инференса на CPU: TensorFlow 2.5.0 в среднем на 5% быстрее, чем TensorFlow 1.15.0 и в среднем на 40% быстрее, чем PyTorch 1.9.0.

Интересно, что при инференсе на CPU переход от динамического графа к статическому в TensorFlow 2.5.0 не влияет на скорость.

4. Количество выделяемой видеопамяти

Для измерения выделяемой видеопамяти я использовал python-оболочку для nvidia-smi. Благодаря этой библиотеке, в конце процесса обучения/инференса я получал информацию о количестве выделенной видеопамяти под текущий процесс. К сожалению, видеопамять в TensorFlow выделяется не совсем понятным образом: как правило, значительная часть доступной видеопамяти выделяется независимо от параметров обучаемой нейронной сети. При этом я несколько раз замечал, что там, где ожидается большее потребление видеопамяти, TensorFlow выделяет её меньше. Например, для обучения ResNet18 с входным размером изображения (128, 128, 3) и размером батча 1 выделяется 4679 Мб видеопамяти, а для обучения той же ResNet18 с входным размером изображения (320, 320, 3) и размером батча 1 выделяется 3911 Мб видеопамяти. В PyTorch выделение видеопамяти более понятно: для описанных конфигураций выделяется 1329 Мб и 1355 Мб видеопамяти соответственно.

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

Таблица 4. Количество выделяемой видеопамяти при обучении и инференсе в PyTorch 1.9.0 (в Мб).

Заключение

Проанализировав результаты измерений, я сделал для себя несколько выводов:

• помимо того, что TensorFlow 1 уже морально устарел, TensorFlow 2 является в целом более производительным при обучении свёрточных нейронных сетей и инференсе без использования GPU;

• использовать tf.keras для инференса на GPU в TensorFlow 2 можно на этапе разработки и тестирования моделей, а для встраивания обученных моделей желательно использовать другие средства TensorFlow, обеспечивающие бoльшую производительность, например, в TensorFlow 2 осталось доступна заморозка графа;

• PyTorch радует своей производительностью при использовании GPU, но если планируется запускать обученные модели на CPU, то необходимо использовать другие средства, например, onnxruntime. На мой взгляд, onnxruntime обеспечивает наилучшую производительность при запуске нейронных сетей на CPU.

Несколько ссылок, где можно подробнее почитать про сравнение различных фреймворков глубокого обучения:

• 1 – небольшое описание и сравнение TensorFlow и PyTorch;

• 2 – сравнение TensorFlow и PyTorch с примерами кода;

• 3 – краткое описание 8 различных фреймворков глубокого обучения.

Весь код, использованный мной для выполнения этих измерений, а также результаты измерений опубликованы в моём репозитории.