Оптимизация моделей Ultralytics YOLO с помощью интеграции TensorRT

Представь беспилотный автомобиль, который едет по оживлённой улице и должен за считаные миллисекунды обнаружить пешехода, ступающего с тротуара на проезжую часть. Одновременно ему может понадобиться распознать знак «Стоп», частично скрытый деревом, или быстро среагировать на соседний автомобиль, резко перестраивающийся в его полосу. В таких ситуациях скорость и реакция в реальном времени критически важны.

Именно здесь ключевую роль играет искусственный интеллект (AI), а конкретно компьютерное зрение — область ИИ, которая помогает машинам интерпретировать визуальные данные. Чтобы решения на основе компьютерного зрения надёжно работали в реальных условиях, им часто нужно быстро обрабатывать информацию, одновременно выполнять несколько задач и эффективно использовать память.

Один из способов добиться этого — аппаратное ускорение с помощью специализированных устройств, например графических процессоров (GPUs), на которых модели работают быстрее. Графические процессоры NVIDIA особенно известны такими задачами благодаря низкой задержке и высокой пропускной способности.

Однако запуск модели на GPU без дополнительной настройки не всегда гарантирует оптимальную производительность. Для полного использования возможностей аппаратных устройств модели ИИ для компьютерного зрения обычно нужно оптимизировать. Чтобы добиться максимальной производительности на конкретном оборудовании, модель необходимо скомпилировать с набором инструкций, предназначенным для этого оборудования.

Например, TensorRT — это формат экспорта и библиотека оптимизации, разработанные NVIDIA для повышения производительности на GPU NVIDIA. В них используются передовые методы, позволяющие значительно сократить время инференса без потери точности.

Рис. 1. NVIDIA TensorRT обеспечивает оптимальную работу моделей на различных устройствах NVIDIA.

В этой статье мы рассмотрим интеграцию с TensorRT, поддерживаемую Ultralytics, и покажем, как экспортировать модель YOLO11 для более быстрого и эффективного развёртывания на оборудовании NVIDIA. Начнём!

Обзор TensorRT#

TensorRT — это набор инструментов, разработанный NVIDIA, чтобы модели ИИ быстрее и эффективнее работали на GPU NVIDIA. Он предназначен для реальных приложений, где скорость и производительность особенно важны, например для беспилотных автомобилей и контроля качества в производстве и фармацевтике.

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

Интеграция TensorRT, поддерживаемая Ultralytics, оптимизирует модель YOLO для более эффективной работы на GPU, используя такие методы, как снижение точности. Это означает, что для представления данных модели используются форматы с меньшим количеством бит, например 16-битные числа с плавающей точкой (FP16) или 8-битные целые числа (INT8). Это сокращает потребление памяти и ускоряет вычисления при минимальном влиянии на точность.

Кроме того, совместимые слои нейронной сети в оптимизированных моделях TensorRT объединяются, чтобы снизить потребление памяти и ускорить инференс.

Рис. 2. Принцип объединения слоёв в TensorRT.

Основные возможности формата экспорта TensorRT#

Прежде чем перейти к экспорту Ultralytics YOLO11 с помощью интеграции TensorRT, рассмотрим основные возможности формата модели TensorRT:

• Простая интеграция с фреймворками: TensorRT поддерживает прямую интеграцию с популярными фреймворками ИИ, такими как PyTorch, Hugging Face и ONNX, обеспечивая производительность до 6 раз выше. Также поддерживается MATLAB, что позволяет разрабатывать высокоскоростные движки ИИ для таких платформ, как Jetson, NVIDIA DRIVE и центры обработки данных.
• Масштабируемое развёртывание с Triton: модели, оптимизированные в формате TensorRT, можно масштабно развёртывать с помощью NVIDIA Triton Inference Server. Эффективность повышают такие возможности, как пакетная обработка входных данных, одновременное выполнение моделей, поддержка ансамблей моделей и потоковая передача аудио и видео в реальном времени.
• Гибкая работа на разных устройствах: TensorRT работает во всей экосистеме NVIDIA — от компактных периферийных устройств до мощных серверов. Он поддерживает такие инструменты, как DeepStream для обработки видео и Riva для речевого ИИ, а также решения для кибербезопасности, рекомендательных систем и других задач.

Как работает интеграция с TensorRT?#

Экспортировать модели Ultralytics YOLO, например Ultralytics YOLO11, в формат модели TensorRT просто. Давай разберём необходимые шаги.

Для начала установи пакет Ultralytics для Python с помощью менеджера пакетов, например pip. Для этого выполни команду pip install ultralytics в командной строке или терминале.

После успешной установки пакета Ultralytics для Python можно обучать, тестировать, дообучать, экспортировать и развёртывать модели для разных задач компьютерного зрения, например обнаружения объектов, классификации и сегментации экземпляров. Если при установке пакета возникнут трудности, решения и советы можно найти в руководстве по распространённым проблемам.

Для следующего шага тебе понадобится устройство NVIDIA. Используй приведённый ниже фрагмент кода, чтобы загрузить и экспортировать YOLO11 в формат модели TensorRT. Код загружает предварительно обученную нано-версию модели YOLO11 (yolo11n.pt) и экспортирует её в файл движка TensorRT (yolo11n.engine), готовый к развёртыванию на устройствах NVIDIA.

from ultralytics import YOLO

model = YOLO("yolo11n.pt")
model.export(format="engine")

Преобразовав модель в формат TensorRT, можно развернуть её для самых разных задач.

В примере ниже показано, как загрузить экспортированную модель YOLO11 (yolo11n.engine) и выполнить с её помощью инференс. Инференс — это использование обученной модели для получения предсказаний на новых данных. В этом примере для проверки модели мы воспользуемся изображением собаки.

tensorrt_model = YOLO("yolo11n.engine")
results = tensorrt_model("https://images.pexels.com/photos/1254140/pexels-photo-1254140.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2.jpg", save=True)

После запуска этого кода выходное изображение будет сохранено в папке runs/detect/predict.

Рис. 3. Результат инференса с использованием экспортированной модели Ultralytics YOLO11 в формате TensorRT.

Когда стоит использовать интеграцию TensorRT#

Пакет Ultralytics для Python поддерживает различные интеграции, позволяющие экспортировать модели YOLO в такие форматы, как TorchScript, CoreML, ONNX и TensorRT. В каких случаях стоит выбрать интеграцию TensorRT?

Вот несколько факторов, отличающих формат модели TensorRT от других вариантов интеграции для экспорта:

• Меньший размер модели: Экспорт модели YOLO в формат TensorRT с точностью INT8 (quantize=8) может значительно уменьшить её размер. Квантование с FP32 до INT8 может сократить размер модели в 4 раза, что ускоряет загрузку, уменьшает требования к хранилищу и снижает объём занимаемой памяти при развёртывании.
• Меньшее энергопотребление: квантование INT8 не только уменьшает размер модели, но и снижает энергопотребление. Операции с пониженной точностью в экспортированных моделях YOLO с INT8 могут потреблять меньше энергии, чем в моделях с FP32. Это особенно полезно для устройств с питанием от батареи, например дронов, смартфонов и периферийных устройств.
• Более высокая производительность: сочетание эффективной архитектуры YOLO с оптимизацией INT8 в TensorRT может повысить скорость инференса.

Применение Ultralytics YOLO11 и формата модели TensorRT#

Модели Ultralytics YOLO, экспортированные в формат TensorRT, можно развернуть в самых разных реальных сценариях. Эти оптимизированные модели особенно полезны там, где важна высокая скорость и эффективность ИИ. Рассмотрим несколько интересных примеров их применения.

Умные кассы самообслуживания в розничных магазинах#

Во многих магазинах сотрудники по-прежнему вручную выполняют самые разные задачи: сканируют штрихкоды, взвешивают товары и упаковывают покупки. Однако зависимость только от сотрудников может замедлять работу и вызывать недовольство покупателей, особенно на кассе. Длинные очереди неудобны и для покупателей, и для владельцев магазинов. Умные кассы самообслуживания — отличное решение этой проблемы.

Такие кассы используют компьютерное зрение и GPU, чтобы ускорить обслуживание и сократить время ожидания. Компьютерное зрение позволяет этим системам видеть и понимать окружающую обстановку, например с помощью обнаружения объектов. Передовые модели, такие как Ultralytics YOLO11, после оптимизации с помощью таких инструментов, как TensorRT, могут работать на GPU намного быстрее.

Экспортированные модели отлично подходят для интеллектуальных систем розничной торговли на компактных, но мощных устройствах, таких как NVIDIA Jetson Nano, разработанный специально для периферийных приложений ИИ.

Рис. 4. Пример интеллектуальной кассы.

Автоматическое обнаружение дефектов в производстве#

Модель компьютерного зрения, например Ultralytics YOLO11, можно обучить на собственных данных для обнаружения дефектной продукции в промышленности. После обучения модель можно экспортировать в формат TensorRT и развернуть на предприятиях с высокопроизводительными системами ИИ.

Когда продукция движется по конвейеру, камеры снимают её, а модель Ultralytics YOLO11 в формате TensorRT анализирует изображения в реальном времени и выявляет дефекты. Это позволяет компаниям быстро и точно обнаруживать проблемы, сокращая количество ошибок и повышая эффективность.

Подобные системы применяются и в фармацевтической промышленности для выявления дефектов медицинской упаковки. По прогнозам, к 2026 году мировой рынок интеллектуальных систем обнаружения дефектов вырастет до 5 миллиардов долларов.

Рис. 5. Обнаружение дефектов в фармацевтической отрасли с помощью YOLO.

Что нужно учитывать при использовании TensorRT#

Хотя интеграция с TensorRT даёт много преимуществ, например ускоряет инференс и снижает задержку, важно помнить о нескольких ограничениях:

• Небольшое снижение точности: После экспорта модели в формат TensorRT точность экспортированной модели может оказаться ниже, чем у исходной. Такие метрики производительности, как точность, полнота и качество обнаружения объектов моделью (значения mAP), могут немного снизиться. Это можно компенсировать, используя репрезентативный калибровочный набор данных (аргумент data) при квантовании INT8.
• Усложнение отладки: Оптимизации TensorRT могут затруднить поиск ошибок и понимание неожиданного поведения, особенно при сравнении результатов с исходной моделью.
• Зависимость от размера батча: Преимущества TensorRT в производительности заметнее при больших размерах батча. Для приложений, обрабатывающих отдельные изображения или небольшие батчи, улучшение производительности может быть менее существенным.

Главные выводы#

Экспорт моделей Ultralytics YOLO в формат TensorRT значительно ускоряет их работу и повышает ее эффективность, что делает их идеальными для задач реального времени: обнаружения дефектов на заводах, работы интеллектуальных касс самообслуживания и наблюдения за оживленными городскими районами.

Эта оптимизация помогает моделям лучше работать на GPU NVIDIA, ускоряя предсказания и снижая потребление памяти и энергии. Несмотря на некоторые ограничения, прирост производительности делает интеграцию с TensorRT отличным выбором для тех, кто создает высокоскоростные системы компьютерного зрения на оборудовании NVIDIA.

Хочешь узнать больше об ИИ? Загляни в наш репозиторий GitHub, присоединяйся к нашему сообществу и изучи варианты лицензирования, чтобы начать свой проект в области компьютерного зрения. Узнай больше о таких инновациях, как ИИ в производстве и компьютерное зрение в логистике, на страницах наших решений.