ONNXRuntime для ускорения ML моделей ⚡️

Сегодня речь пойдет про формат ONNX. Он помогает портировать ML модели в стандартном формате на другие языки, например на JVM. Для запуска моделей в этом формате используется фреймворк ONNXRuntime. Кроме того ONNXRuntime позволяет ускорить инференс. Об этом мы сегодня и поговорим. В качестве примера возьмем модель на PyTorch и сравним скорость инференса до и после конвертация в ONNX.

Как работает ускорение?

В PyTorch обработка данных ленивая, т.е. команды сначала добавляются в очередь и потом из нее выполняются. Для выполнения всех запланированных вычислений нужен torch.cuda.synchronize(). Запустим torch модель на входном массиве (1,3,224,224) , запишем время инференса и оценим FPS. Также удалим, из результатов первые 20 запусков, которые как правило уходят на прогрев.

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

torch.cuda.synchronize()

for step in range(nsteps):
 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))

PyTorch выдает 400 FPS. Этот результат получился в следствие того, что граф модели PyTorch - динамический, что ограничивает его скорость инференса. Теперь скастуем из динамического графа torch'а статический граф и запишем в формат ONNX

torch.onnx.export(
 model,
 torch.randn(1, 224, 224).to(device),
 'mobilenet_v2.onnx',
 input_names = ['input'],
 output_names = ['output'],   
)

Для того, чтобы подгрузить ONNX модель, нам нужно создать InferenceSession. Далее читаем модель и запускаем инференс на ONNXRuntime. Также засечем время инференса, чтобы проверить ускорение.

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, т.е. мы получаем прирост скорости инференса в полтора раза. Здесь запуск проводился на CUDA, но ONNX также поддерживает и запуски на CPU