Автор: Андрей Щербин Редактура: Александр Гончаренко
Подход «на чём обучали, на том и запускайте» — не самое эффективное решение с точки зрения быстрого действия и использования вычислительных ресурсов.
Граф модели в PyTorch — динамический, что ограничивает скорость его расчёта, поскольку в динамическом графе исполнение CUDA-кернела происходит во время вызова операции, связанной с этим кернелом. Из-за этого получаются накладные расходы на запуск. Кроме того, многие операции, которые нужны при тренировке, во время инференса могут быть удалены.
Например, если BatchNorm идёт после свёртки — их можно «сплавить» (fuse) в одну операцию, т.к. эти преобразования линейные. Подобного рода оптимизации могут быть применены также к операциям reshape, squeeze, unsqueeze.
Для запуска оптимизаций сначала нужно получить статический граф вычислений. Это можно сделать, если сохранить модель в формате ONNX (Open Neural Network eXchange). Как видно из названия, он позволяет делиться моделями между разными фреймворками. Но у Microsoft, создателя этого формата, есть также свой фреймворк инференса – ONNXRuntime. Он ****не только помогает запускать модели на CPU и GPU более эффективно, чем обычный PyTorch (ведь граф теперь статический), но и поддерживает смену execution provider'а на более низкоуровневый (согласно документации есть даже и для GPU от AMD).
Вообще, к запуску на целевом железе обученной модели стоит подходить как к отдельной задаче. Часто целевым устройством становится процессор из-за отсутствия бюджета на видеокарты. Но многие забывают, что в нём бывает встроенная GPU, которую тоже можно использовать для инференса нейросетей. Запуск модели на ней доступен, например, из фреймворка OpenVINO.
Если же вы используете в качестве вычислителя видеокарту от Nvidia, напрашивается запускать её на TensorRT. Или если вы понимаете, какие именно вычисления происходят в вашей модели, вы можете просто закодить их на CUDA. Узнать о том, как это делается, можно с помощью этого доклада. TensotRT и особенности его работы мы разберём чуть позже, но сейчас отметим — он использует более эффективные CUDA-ядра, чем PyTorch.
Представим следующую ситуацию: ваша компания вместо устаревающих 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))