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

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

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

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

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

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

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

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

Обзор TensorRT#

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

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

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

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

Рис. 2. Обзор метода объединения слоёв в TensorRT.

Ключевые особенности формата экспорта TensorRT#

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

• Простая интеграция с фреймворками: TensorRT поддерживает прямую интеграцию с популярными фреймворками AI, такими как PyTorch, Hugging Face и ONNX, обеспечивая производительность до 6 раз выше. Он также поддерживает MATLAB, позволяя разрабатывать высокоскоростные AI-движки на таких платформах, как Jetson, NVIDIA DRIVE и центры обработки данных.

• Масштабируемое развёртывание с Triton: Модели, оптимизированные в формате TensorRT, можно масштабно развёртывать с помощью NVIDIA Triton Inference Server, который повышает эффективность благодаря пакетной обработке входных данных, одновременному выполнению моделей, поддержке ансамблей моделей и потоковой передаче аудио и видео в реальном времени.

• Гибкость на разных устройствах: От небольших периферийных устройств до мощных серверов TensorRT работает во всей экосистеме NVIDIA, поддерживая такие инструменты, как DeepStream для видео, Riva для речевого AI, а также решения для кибербезопасности, рекомендаций и других задач.

Как работает интеграция 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 может значительно уменьшить размер модели. Квантизация с FP32 до INT8 может сократить размер модели в 4 раза, что ускоряет загрузку, снижает требования к хранилищу и уменьшает объём памяти, необходимый при развёртывании.

• Меньшее энергопотребление: Квантизация INT8 не только уменьшает размер модели, но и снижает энергопотребление. Операции с пониженной точностью для экспортированных в INT8 моделей YOLO могут потреблять меньше энергии по сравнению с моделями FP32, что особенно полезно для устройств с питанием от батареи, таких как дроны, смартфоны и периферийные устройства.

• Более высокая производительность: Сочетание эффективной архитектуры YOLO с оптимизацией INT8 в TensorRT может повысить скорость инференса.

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

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

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

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

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

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

Рис. 4. Пример умной кассы самообслуживания.

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

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

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

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

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

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

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

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

• Повышенная сложность отладки: Оптимизации, выполняемые TensorRT, могут усложнить трассировку ошибок или понимание неожиданного поведения, особенно при сравнении результатов с исходной моделью.

• Зависимость от размера пакета: Преимущества TensorRT в производительности более заметны при больших размерах пакета. Для приложений, обрабатывающих отдельные изображения или небольшие пакеты, улучшение производительности может быть менее существенным.

Основные выводы#

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

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

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