⁣Zero-Copy инференс ONNX в Python: убиваем лишние копирования

При развертывании ONNX-моделей в production каждая микросекунда на счету. Часто узким местом становится не сам инференс, а оверхэд на передачу данных между Python и C++ рантаймом ONNX Runtime. Решение — zero-copy тензоры.

По умолчанию ONNX Runtime (ORT) делает копию входных тензоров из numpy-массивов в свой внутренний формат. Для больших батчей это может добавить 10-30% времени к инференсу. Я наступал на эти грабли: модель считает 5 мс, а общее время запроса — 7 мс, и ты гадаешь, куда утекает.

Как работает zero-copy с OrtValue

Вместо ort_session.run(None, {'input': numpy_array}) можно передавать OrtValue напрямую:

import onnxruntime as ort
import numpy as np

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
ort_input = ort.OrtValue.ort_value_from_numpy(input_data)

session = ort.InferenceSession('model.onnx')
results = session.run_with_ort_values({'input': ort_input})
output_tensor = results[0].numpy()

Тут есть нюанс: run_with_ort_values возвращает OrtValue, и .numpy() может вернуть view на те же данные с нулевым копированием, если тензор выровнен.

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

Чтобы минимизировать оверхэд еще сильнее, можно выделить буфер один раз и переиспользовать его:

buffer_np = np.empty((1, 3, 224, 224), dtype=np.float32)
buffer_ort = ort.OrtValue.ort_value_from_numpy(buffer_np)

# В цикле инференса
np.copyto(buffer_np, new_input_data)
results = session.run_with_ort_values({'input': buffer_ort})

Замерил на batch_size=1, 1000 запросов, ONNX Runtime 1.18:
- Стандартный run: 2.1 мс на запрос (включая 0.3 мс на копирование)
- Zero-copy с OrtValue: 1.8 мс (копирование 0 мс)
- + предварительная аллокация: 1.7 мс

Zero-copy дает около 15% ускорения. Для высоконагруженных сервисов это уже заметно.

Типичная ошибка: когда zero-copy не сработает

Если входной тензор не выровнен (non-contiguous) — ORT сделает copy. Спасает np.ascontiguousarray(). Если модель меняет форму тензора на входе — копирование неизбежно, тут ничего не поделать.

Рекомендации для прода

- Используйте OrtValue.ort_value_from_numpy() вместо run().
- Переиспользуйте OrtValue-буферы.
- Включайте ort.SessionOptions().enable_cpu_mem_arena = True.

Не ждите чуда, но 15% выньете просто так.

Вывод:
Zero-copy с OrtValue и переиспользование буферов — простой и надежный способ снизить latency инференса ONNX-моделей на 10-15% без изменения архитектуры.