TensorRT

TensorRT — это программная библиотека (SDK) для оптимизации и ускорения выполнения нейросетевых моделей на графических процессорах (GPU) производства компании NVIDIA. Библиотека предназначена для инференса (вывода) обученных моделей в промышленных, встраиваемых и облачных системах, обеспечивая максимальную производительность за счёт низкоуровневой оптимизации под архитектуру CUDA и тензорных ядер (Tensor Cores) GPU.

История

Разработка TensorRT началась компанией NVIDIA в середине 2010-х годов как ответ на растущую потребность в эффективном запуске глубоких нейронных сетей на GPU в реальном времени. Первая публичная версия, известная как NVIDIA TensorRT 1.0, была выпущена в 2016 году. Изначально библиотека позиционировалась как средство для ускорения инференса в центрах обработки данных и на высокопроизводительных графических ускорителях (например, Tesla P100).

Ключевым этапом развития стало внедрение поддержки тензорных ядер (Tensor Cores) в архитектуре Volta (2017 год). TensorRT получил возможность автоматически преобразовывать операции с плавающей запятой одинарной точности (FP32) в смешанную точность (FP16/INT8), что дало многократный прирост производительности без существенной потери точности. В последующих версиях (TensorRT 5, 6, 7, 8) были добавлены поддержка новых архитектур GPU (Turing, Ampere, Hopper), расширены возможности калибровки для INT8 и INT4, а также улучшена интеграция с популярными фреймворками глубокого обучения.

В 2023 году вышла версия TensorRT 8.6, которая представила поддержку моделей на основе Transformer-архитектуры (включая GPT, BERT) и улучшила работу с динамическими размерами входных данных. В 2024 году была выпущена TensorRT 10.0, в которой была значительно переработана архитектура парсера моделей, а также добавлена экспериментальная поддержка разрежённых вычислений (sparsity).

Архитектура и принцип работы

TensorRT работает в два основных этапа: построение оптимизированного плана (build) и выполнение инференса (runtime).

Этап построения (Build Phase)

На этом этапе TensorRT принимает на вход обученную модель в одном из поддерживаемых форматов (ONNX, TensorFlow SavedModel, Caffe prototxt, PyTorch через экспорт в ONNX) и выполняет серию оптимизаций:

1. Развёртывание графа (Graph optimization): Исходный вычислительный граф модели анализируется и преобразуется. Устраняются избыточные операции (например, последовательные слои Batch Normalization и ReLU сливаются в один слой CBR — Convolution + BatchNorm + ReLU).
2. Слияние слоёв (Layer fusion): Несколько последовательных слоёв объединяются в один оптимизированный слой (kernel), что уменьшает накладные расходы на вызов ядер CUDA.
3. Выбор формата данных: TensorRT выбирает наиболее эффективный формат данных для каждого слоя: FP32, FP16, INT8 или INT4. Для INT8 требуется калибровка — подача репрезентативного набора входных данных для определения диапазона активаций.
4. Планирование памяти: Оптимизируется распределение памяти GPU для минимизации фрагментации и ускорения доступа.
5. Генерация плана (Plan file): Результатом является бинарный файл (расширение .plan), который содержит оптимизированный граф, веса модели и метаданные для выполнения.

Этап выполнения (Runtime Phase)

На этапе выполнения TensorRT загружает план в память GPU и выполняет инференс. Runtime-компонент библиотеки минимален по размеру (около 1–2 МБ) и не требует полной установки TensorRT, что удобно для развёртывания на конечных устройствах.

Поддерживаемые типы данных и точность

TensorRT поддерживает несколько форматов данных, что позволяет балансировать между производительностью и точностью:

Интеграция с фреймворками

TensorRT интегрируется с основными фреймворками глубокого обучения через специализированные обёртки:

• TensorFlow-TensorRT (TF-TRT): Позволяет оптимизировать модели TensorFlow непосредственно в процессе обучения или после него. Поддерживает автоматическую замену подграфов на оптимизированные ядра TensorRT.
• PyTorch-TensorRT: Модуль torch2trt или официальный torch_tensorrt (начиная с PyTorch 1.9) позволяет конвертировать модели PyTorch в TensorRT планы.
• ONNX Runtime: Поддерживает бэкенд TensorRT для выполнения ONNX-моделей на GPU.
• NVIDIA Triton Inference Server: Сервер инференса, который использует TensorRT как один из бэкендов для высокопроизводительного развёртывания.

Применение

TensorRT широко используется в системах, где требуется низкая задержка (latency) и высокая пропускная способность (throughput) при выполнении нейросетей:

• Автономные транспортные средства: Обработка данных с камер и лидаров в реальном времени (например, в платформе NVIDIA DRIVE).
• Робототехника: Управление роботами на основе компьютерного зрения.
• Облачные сервисы: Развёртывание моделей рекомендательных систем, обработки естественного языка (NLP) и компьютерного зрения в дата-центрах (например, в сервисах Amazon SageMaker, Google Cloud AI Platform).
• Встраиваемые системы: Запуск моделей на устройствах NVIDIA Jetson (TX2, Xavier, Orin) для периферийных вычислений (edge AI).
• Медицинская диагностика: Ускорение анализа медицинских изображений (МРТ, КТ, рентген).

Ограничения и критика

Несмотря на высокую эффективность, TensorRT имеет ряд ограничений:

1. Привязка к оборудованию: Оптимизированные планы привязаны к конкретной архитектуре GPU и версии драйвера. План, созданный для GPU Ampere, не будет работать на GPU Volta.
2. Сложность отладки: Внутренние оптимизации (слияние слоёв, изменение формата) делают отладку модели сложной — разработчик не может легко просмотреть промежуточные активации.
3. Ограниченная поддержка операций: Не все операции, доступные в популярных фреймворках, поддерживаются TensorRT. В случае неподдерживаемой операции модель либо не конвертируется, либо падает на этапе выполнения.
4. Необходимость калибровки: Для использования INT8 требуется тщательная калибровка на репрезентативном наборе данных, иначе точность модели может существенно снизиться.
5. Закрытый исходный код: В отличие от альтернатив (например, Apache TVM или OpenVINO), TensorRT является проприетарным продуктом NVIDIA, что ограничивает возможности модификации и аудита.

Сравнение с альтернативами

Интересные факты

• TensorRT используется в NVIDIA GeForce NOW для стриминга игр с ИИ-улучшением изображения (DLSS).
• В 2022 году компания NVIDIA продемонстрировала ускорение инференса модели GPT-3 (175 млрд параметров) с использованием TensorRT и технологии Megatron-LM, достигнув пропускной способности более 10 000 токенов в секунду на одном сервере с 8 GPU A100.
• Версия TensorRT 10.0 (2024) ввела поддержку так называемых «плагинов» — пользовательских C++/CUDA-ядер, которые могут быть встроены в оптимизированный граф.

Источники

1. NVIDIA Developer Documentation: TensorRT Developer Guide (версия 10.0, 2024).
2. NVIDIA Technical Blog: «Accelerating Inference with TensorRT» (2017–2024).
3. Документация TensorFlow-TensorRT (TF-TRT) на GitHub.
4. Документация PyTorch-TensorRT (torch_tensorrt) на GitHub.
5. Статья «A Survey of Deep Learning Inference Accelerators» (IEEE Access, 2023).