«PyTorch is all you need»?

Введение

Подход «на чём обучали, на том и запускайте»‎ — не самое эффективное решение с точки зрения быстрого действия и использования вычислительных ресурсов.

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

Например, если BatchNorm идёт после свёртки — их можно «сплавить» (fuse) в одну операцию, т.к. эти преобразования линейные. Подобного рода оптимизации могут быть применены также к операциям reshape, squeeze, unsqueeze.

Альтернативы PyTorch

Для запуска оптимизаций сначала нужно получить статический граф вычислений. Это можно сделать, если сохранить модель в формате ONNX (Open Neural Network eXchange). Как видно из названия, он позволяет делиться моделями между разными фреймворками. Но у Microsoft, создателя этого формата, есть также свой фреймворк инференса – ONNXRuntime. Он ****не только помогает запускать модели на CPU и GPU более эффективно, чем обычный PyTorch (ведь граф теперь статический), но и поддерживает смену execution provider’а на более низкоуровневый (согласно документации есть даже и для GPU от AMD).

Вообще, к запуску на целевом железе обученной модели стоит подходить как к отдельной задаче. Часто целевым устройством становится процессор из-за отсутствия бюджета на видеокарты. Но многие забывают, что в нём бывает встроенная GPU, которую тоже можно использовать для инференса нейросетей. Запуск модели на ней доступен, например, из фреймворка OpenVINO.

Если же вы используете в качестве вычислителя видеокарту от Nvidia, напрашивается запускать её на TensorRT. Или если вы понимаете, какие именно вычисления происходят в вашей модели, вы можете просто закодить их на CUDA. Узнать о том, как это делается, можно с помощью этого доклада. TensotRT и особенности его работы мы разберём чуть позже, но сейчас отметим — он использует более эффективные CUDA-ядра, чем PyTorch.

ONNX и ONNXRuntime

Представим следующую ситуацию: ваша компания вместо устаревающих 1080ti закупила для обучения 4090, а старые карты решила использовать для серверов инференса. Для простоты положим, что на одном сервере используется одна видеокарта (но это можно масштабировать). Итак, вам дают доступ к такому серверу вместе с запущенным на нём эмбеддером лиц на основе MobileNetV2, а ещё — код инференса на PyTorch, написанный стажёром (по факту копипаст валидации, но без подсчёта метрик) по принципу «и так сойдёт».

Решение задачи оптимизации инференса на конкретном устройстве обычно начинается с профилирования baseline-решения. Бывает, что основное время тратится не на вычисление нейронной сети, а на препроцессинг / постпроцессинг данных. Это случается, если препроцессинг выполняется на процессоре, а процессор в целевом устройстве значительно слабее видеочипа. Пример такого устройства — Nvidia Jetson. Если вычисление нейронной сети — самый долгий этап вашего пайплайна, а ваша задача — его ускорение (остальным пайплайном займутся коллеги), то необходимо замерить время вычисления нейронной сети в baseline-решении.

Чтобы замерить latency модели на PyTorch, важно выполнить 10-20 предварительных запусков для прогрева. Он нужен для моделирования реальных условий работы системы: на практике запросы часто приходят в систему непрерывно, и при старте обработки очередного запроса все нужные библиотеки и CUDA kernel’ы подгружаются в память.

При замере времени нужно помнить, что вычисления на CUDA в PyTorch ленивые, как и отправка данных. Следовательно, при исполнении строки кода, которая подразумевает отправку данных на устройство с CUDA или запуск какой-то операции реальной отправки, вычисления не происходят. Эти команды просто добавляются в очередь, а их реальное исполнение может случиться в произвольный момент времени. Для выполнения всех запланированных отправок / вычислений необходимо выполнить torch.cuda.synchronize(). Тогда замер времени вычисления модели будет выглядеть следующим образом:

n_steps = 200
n_warmup_steps = 20
times = []
data = torch.ones((1,3,224,224), device='cuda')

torch.cuda.synchronize()

for step in range(n_steps):
		start_inference = time.time()
		result = model(data)
		torch.cuda.synchronize()
		end_inference = time.time()
		times.append(end_inference - start_inference)
times = times[n_warmup_steps:]
print(np.mean(times), np.std(times))

Итак, теперь наш эмбеддер при размере входа (1, 3, 224, 224) на видеокарте Nvidia 1080ti с использованием PyTorch способен выдавать около 400 FPS (то есть обрабатывать 400 запросов в секунду). Выглядит неплохо! Но компания хочет выжать максимум из сервера, чтобы больше клиентов смогло воспользоваться вычислительными ресурсами без покупки новых.

Сначала давайте избавимся от динамического графа за счёт конвертации модели в формат ONNX и запустим её на ONNXRuntime. Здесь исполнение неленивое, а значит, никакой синхронизации. Конвертация PyTorch-модели в ONNX выполняется при помощи следующего кода:

torch.onnx.export(
		model,                                # model being run
		torch.randn(1, 224, 224).to(device),  # model input (or a tuple for multiple inputs)
		'mobilenet_v2.onnx',                  # where to save the model (can be a file or file-like object)
		input_names = ['input'],              # the model's input names
		output_names = ['output'],            # the model's output names
)

Код замера времени исполнения будет выглядеть следующим образом:`

import onnxruntime as ort
import numpy as np

providers = [('CUDAExecutionProvider', {'device_id': 0})]

ort_sess = ort.InferenceSession(‘mobilenet_v2.onnx', providers=providers)

n_steps = 200
n_warmup_steps = 20

times = []

image_as_numpy = np.ones((1,3,224,224))

for step in range(n_steps):
	start_inference = time.time()
	outputs = ort_sess.run(None, {'input': image_as_numpy})
	end_inference = time.time()
	times.append(end_inference - start_inference)

times = times[n_warmup_steps:]
print(np.mean(times), np.std(times))

При запуске модели через ONNXRuntime получаем 630 FPS — ускорение в 1.6 раза за счёт статического графа вычислений и конвертации в ONNX.

Кроме CUDA ONNXRuntime поддерживает запуск на CPU и более низкоуровневых фреймворках (как Apache TVM), популярный NPU, а также OpenVINO для ускорения вычислений на процессоре или встроенных в него видеокартах Intel.

Также ONNXRuntime поддерживает фреймфорк XNNPACK для запуска моделей на Android и IOS. Есть и провайдер для TensorRT, который позволяет использовать все преимущества низкоуровневых фреймворков инференса из одного интерфейса ONNXRuntime и не переписывать код инференса под различные hardware платформы. Более подробную информацию о провайдерах можно получить в официальной документации.

Возникает вопрос: можно ли добиться большего ускорения? Ответ: можно! Для этого необходимо использовать более низкоуровневый и приближенный к железу фреймворк — TensorRT.

Осторожно, спойлер: MobileNetV2 на GTX 1080ti при размере входа (1,3,224,224) способен выдавать 1200 FPS на TensorRT.

TensorRT

Что же такое TensorRT, и почему он лучше обеспечивает ускорение в сравнении с ONNXRuntime?

В нём модель оптимизируется перед запуском таким образом, чтобы для вычисления каждой операции использовались самые быстрые алгоритмы на конкретном железе с определённым количеством свободной памяти (именно в целевом приложении).

Перед отправкой модели в TensorRT нужно конвертировать её в ONNX — это удобный способ передачи модели между фреймворками. TensorRT поддерживает и динамическое построение модели через создание NetworkDefinition. Некоторые фреймворки (например, torch2trt) используют такую функцию для вычисления в TensorRT частей модели (другие вычисляются на PyTorch). Но обычно это медленнее и связано с наличием у модели операций, которые не поддерживаются TensorRT или ONNX.

Полезная вещь для ONNX-модели — применение оптимизаций из onnx-simplifier или onnx-graphsurgeon. PyTorch конвертирует модель в ONNX операция за операцией, а эти библиотеки распознают их паттерны, заменяют на более простые аналоги, не нарушая семантику вычислительного графа. Пример оптимизации с помощью onnx-simplifier представлен на картинке ниже. Перед запуском TensorRT для упрощения ему работы рекомендуется применить onnx-simplifier к вашему целевому onnx.

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

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

При оптимизации модели TensorRT заменяет операции в ONNX на соответствующие CUDA-ядра. Схематично это выглядит следующим образом:

На схеме видно, что произошло не только сплавление операций, но и замена трёх веток исполнения на одну. Также при оптимизации можно изменить тип данных, используемый для вычислений. Стандарт здесь — применение FP16 вместо обычного FP32. Однако если быстродействие в FP16 вас не устраивает, и на вашем устройстве присутствуют блоки для целочисленных вычислений, можно квантовать с помощью средств TensorRT.

trtexec

Для быстрой проверки ускорения модели без кода используется утилита командной строки trtexec. Она оценивает latency / throughput с различными опциями оптимизации, а также сохраняет собранный engine с возможностью дальнейшей загрузки и запуска. Она содержит много флагов, которые включают определённые опции оптимизации. О них можно более подробно почитать в документации к утилите, мы же разберём в качестве примера её использование победителями конкурса LPCV2023. Команда выглядит так:

trtexec --workspace=4096 --onnx=model.onnx --saveEngine=engine.trt --best --useSpinWait --outputIOFormats=fp16:chw --inputIOFormats=fp16:chw --verbose

Первый аргумент задаёт максимальный объём используемой памяти в мегабайтах, то есть для вычисления модели будет доступно 4 гигабайта видеопамяти.

Следующие два аргумента указывают пути до модели в формате ONNX и до собранного engine (один из результатов запуска утилиты).

Флаг --best разрешает использовать любые типы данных для достижения максимальной производительности. Дело в том, что в TensorRT применение типов данных носит рекомендательный характер, а финальное решение о том, использовать ли, например, FP16 принимается оптимизатором. На нашей практике флаг --best заставляет оптимизатор чаще выбирать низкоразрядные типы данных.

Флаг --useSpinWait включает активную синхронизацию между GPU и CPU, что увеличивает нагрузку на процессор, но ускоряет синхронизацию. Его использование уменьшает latency даже на устройствах с маломощными процессорами по типу Nvidia Jetson Nano.

Следующие флаги задают порядок осей и тип входных / выходных данных модели. В конкурсе LPCV2023 решалась задача сегментации — выход модели имел сравнимый со входом размер. Применение FP16 для входа и выхода позволило сократить время на передачу изображений в модель и получение результата из неё.

Вернёмся к нашему эмбеддеру на MobileNetV2 и 1080ti. Запустим его через trtexec:

trtexec --workspace=8192 --onnx=mnv2.onnx --saveEngine=mnv2.trt --best --verbose

В итоге при использовании TensorRT мы получаем 1100 FPS для этой модели, что в 2.7 раза больше изначальных результатов на PyTorch. А значит, экономим для нашей компании на вычислительных ресурсах кучу денег — можно идти просить прибавку к зарплате 😄.

Заключение

Возможно, у вас возник вопрос: а зачем возиться с TensorRT, если есть ONNXRuntime с провайдером TensorRT? И вы правы, если задумываетесь о скорости вычисления модели: выигрыш от TensorRT будет, но незначительный.

Однако часто приходится беспокоиться о скорости пайплайна целиком: от декодирования изображения до результата после постпроцессинга. В таком случае хорошим тоном будет применение фреймворка Triton. Он построен поверх TensorRT и несовместим с ONNXRuntime. Поэтому TRT — надёжный способ инференса, особенности работы которого важно знать каждому профессионалу 😉.