Использование TensorRT и ONNX для ускорения инференса моделей глубокого обучения, особенно в задачах компьютерного зрения (CV), стало критически важным компонентом пайплайнов в индустрии, где требуется минимальная задержка и высокая производительность. Ниже описан подробный опыт, методы, типичные сценарии и технические детали применения TensorRT и ONNX.

1. ONNX: Open Neural Network Exchange

Назначение:

ONNX — это открытый формат для представления нейросетевых моделей, разработанный Microsoft и Facebook, позволяющий переносить модели между фреймворками (PyTorch, TensorFlow, Keras и др.) и запускать их на универсальных оптимизированных инференс-движках (ONNX Runtime, TensorRT, OpenVINO).

Поддерживаемые задачи:

• Классификация

• Детекция (YOLO, SSD, Faster R-CNN)

• Сегментация (U-Net, DeepLab)

• OCR (CRNN, TrOCR)

• Pose Estimation

Примеры конвертации:

import torch
model = MyModel()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=11, input_names=\["input"\], output_names=\["output"\])

Основные возможности:

• Кросс-фреймворковая совместимость

• Оптимизация графа (fuse, prune, quantize)

• Работа на CPU, GPU, TPU, FPGA через ONNX Runtime

ONNX Runtime:

• Поддержка Execution Providers: CUDA, TensorRT, OpenVINO, DML, NPU

• Возможность Dynamic Quantization, Int8 calibration, Sparse attention

• Используется как backend в HuggingFace, Microsoft Azure AI

2. TensorRT: NVIDIA TensorRT Inference Engine

Назначение:

TensorRT — фреймворк от NVIDIA для высокоэффективного инференса на GPU с поддержкой квантизации, фьюзинга, оптимизации графа, reordering и dynamic batching. Поддерживает работу напрямую с .onnx, .engine, .plan файлами.

Поддерживаемые задачи:

• CV: YOLOv5/v8, SSD, UNet, CenterNet

• NLP: BERT, GPT, T5 (через FasterTransformer)

• Multimodal: CLIP, SAM (в упрощенной форме)

• Custom DNN: свои модели (через плагины)

Поддерживаемые типы данных:

• FP32 (обычный float)

• FP16 (вдвое быстрее, меньше памяти)

• INT8 (квантизация, требует калибровки)

• Sparsity-aware (структурное разреживание)

3. Пайплайн ONNX → TensorRT

1. Обучение модели в PyTorch / TensorFlow

2. Экспорт в ONNX (torch.onnx.export, tf2onnx)

3. Проверка модели через onnxruntime (в том числе onnxsim, onnxchecker)

4. Оптимизация ONNX (onnxoptimizer, onnx-simplifier)

5. Конвертация в TensorRT Engine:

trtexec --onnx=model.onnx --saveEngine=model_fp16.engine --fp16

1. Использование модели:

import tensorrt as trt
\# Инициализация, загрузка, обработка входов

4. Опыт оптимизации конкретных моделей

YOLOv5 / YOLOv8:

• Экспорт модели в ONNX (с opset_version ≥ 12)

• ONNX simplification + кастомная обработка для Grid/Concat/Sigmoid

• Использование trtexec или Python API (tensorrt.Builder)

• Достижение ускорения до 3–4× при INT8 inference на Jetson Xavier и A100

UNet / DeepLabV3+:

• Проблемы с динамической размерностью маски — решаются через статическое pad/crop

• Часто требуется ручная вставка Resize и Cast-операций

• Сильно выигрывает от FP16 (в 2× быстрее без падения точности)

BERT-like модели:

• Использование TensorRT-plugins через NVIDIA FasterTransformer

• Поддержка batch_size=1 и динамического input length

• Высокая чувствительность к precision (иногда INT8 вызывает degradation)

OCR (CRNN, TrOCR):

• Обработка seq2seq моделей требует упрощения декодера или beam-search как постобработки на CPU

• Часто split на encoder → decoder для частичной оптимизации

5. Сравнение скорости (опытные значения)

Оценки на NVIDIA RTX 3080, INT8 требует calibration dataset.

6. Квантизация (INT8)

Пост-тренировочная квантизация (PTQ):

• Использование calibration data (100–1000 изображений)

• ONNX: quantize_dynamic, quantize_static (onnxruntime-tools)

• TensorRT: INT8 calibrator class

QAT (Quantization Aware Training):

• Обучение с имитацией квантизации

• Требует PyTorch QAT-aware модели

• Лучшие результаты по точности

7. Проблемы и обходные пути

8. Edge-инференс с TensorRT

Jetson Nano / Xavier:

• Использование JetPack SDK (включает TensorRT)

• Прямая конвертация YOLOv5 в .engine

• Использование cv2.dnn или pycuda + tensorrt

Coral, Raspberry Pi (альтернатива):

• Использование TFLite с INT8 квантизацией

• Поддержка ONNX ограничена, TensorRT не работает напрямую

9. Интеграция в продакшен

• Использование Triton Inference Server (поддержка ONNX и TensorRT)

• REST/gRPC интерфейс

• Адаптивный batching, multi-model deployment

• Мониторинг latency / throughput / memory

10. Профилирование и отладка

• trtexec --loadEngine=model.engine --dumpProfile

• nsight-systems, nvprof — профилировка по слоям

• onnxruntime.InferenceSession(..., providers=["TensorrtExecutionProvider"])

Использование ONNX и TensorRT позволяет сократить latency, увеличить throughput и уменьшить энергопотребление модели, что особенно важно в edge-приложениях, мобильных и серверных системах реального времени. Успешное применение требует глубокой интеграции в ML-пайплайн, учёта особенностей каждой архитектуры и внимательного профилирования.