Обучение YOLO: полное руководство по настройке, датасетам и оптимизации

Разбираемся, как обучить нейросеть YOLO под свои задачи: от подготовки датасета и выбора версии до настройки гиперпараметров и запуска на GPU.

Что такое YOLO и почему это важно

YOLO (You Only Look Once) — это семейство нейросетей для компьютерного зрения, которые решают задачу обнаружения объектов за один проход. В отличие от более ранних подходов, которые сначала искали кандидатов на объекты, а затем уточняли их положение, YOLO одновременно предсказывает координаты рамок и классы объектов на изображении. Это делает её невероятно быстрой и позволяет работать в реальном времени даже на относительно простом оборудовании.

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

Семейство YOLO постоянно развивается: каждая новая версия приносит улучшения в скорости, точности и удобстве использования. Сегодня доступны такие версии, как YOLOv5, YOLOv8, YOLOv11 и YOLOv12, и выбор конкретной модели зависит от задачи, доступного оборудования и требований к производительности.

Основные задачи, которые решает YOLO

YOLO — это не только детекция объектов. Современные версии поддерживают несколько задач компьютерного зрения:

  • Детекция — обнаружение объектов и их классов на изображении или видео.
  • Сегментация — разделение изображения на области, относящиеся к каждому объекту, что позволяет, например, удалить фон.
  • Классификация — определение того, что находится на изображении в целом.
  • Поиск ключевых точек — определение позы человека по ключевым точкам тела.
  • Трекинг объектов — потоковая обработка, при которой для каждого объекта сохраняется история местоположения, что позволяет предсказывать движение.

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

Выбор версии YOLO: v5, v8, v11 или v12

На сегодняшний день наиболее популярны YOLOv5 и YOLOv8, а также более новые YOLOv11 и YOLOv12. Каждая версия имеет свои особенности:

  • YOLOv5 — классика, которая завоевала популярность благодаря простоте использования и хорошей документации. Подходит для большинства задач, но уступает новым версиям в точности.
  • YOLOv8 — значительный шаг вперёд: улучшенная архитектура, поддержка сегментации и трекинга, а также удобный API через библиотеку ultralytics. Это самая распространённая версия для обучения под свои данные.
  • YOLOv11 — новое поколение с улучшенной точностью и скоростью, но требует более мощного оборудования для обучения.
  • YOLOv12 — самая свежая версия, которая продолжает развивать идеи предшественников, но пока менее распространена в продакшене.

При выборе версии учитывайте не только точность, но и совместимость с вашим оборудованием. Если у вас ограниченные ресурсы, YOLOv8n (nano) — отличный старт, так как она быстро работает и легко обучается. Для сложных задач, где важна максимальная точность, стоит рассмотреть YOLOv11 или YOLOv12 с большими вариантами моделей.

Подготовка собственного датасета

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

  1. Сбор изображений — соберите не менее 100–200 изображений на каждый класс, которые вы хотите детектировать. Изображения должны быть разнообразными: разные ракурсы, освещение, фон, размеры объектов.
  2. Разметка — для детекции необходимо разметить каждый объект на изображении, указав координаты ограничивающей рамки и класс. Для этого используются инструменты вроде LabelImg, CVAT или Roboflow.
  3. Формат данных — YOLO использует текстовые файлы с координатами в нормализованном виде: class x_center y_center width height. Каждое изображение сопровождается файлом .txt с таким же именем.
  4. Разделение на выборки — разделите датасет на обучающую (train), валидационную (val) и тестовую (test) выборки в пропорции примерно 70/20/10.
  5. Структура папок — создайте структуру, которую ожидает YOLO: images/train, images/val, labels/train, labels/val.

После подготовки датасета создайте файл конфигурации data.yaml, в котором укажите пути к папкам, количество классов и их названия. Этот файл используется при обучении.

Установка и запуск обучения

Для работы с YOLO проще всего использовать библиотеку ultralytics, которая поддерживает YOLOv8 и YOLOv11. Установка выполняется одной командой:

pip install ultralytics

После установки можно начать обучение на своём датасете. Базовый скрипт выглядит так:

from ultralytics import YOLO

# Загрузка предобученной модели
model = YOLO('yolov8n.pt')

# Обучение на своём датасете
model.train(data='data.yaml', epochs=100, imgsz=640)

При первом запуске модель автоматически скачает предобученные веса. Обучение можно запустить как на GPU, так и на CPU, но для больших датасетов GPU обязателен. После завершения обучения модель сохранится в папке runs/detect/train, и вы сможете использовать её для инференса:

model = YOLO('runs/detect/train/weights/best.pt')
results = model('test.jpg')

Важно: не забывайте указывать параметр epochs — количество эпох обучения. Рекомендуемое значение — 300, но для небольших датасетов может хватить и 100.

Настройка гиперпараметров для лучшей точности

Качество обучения YOLO сильно зависит от гиперпараметров. Вот ключевые из них:

  • epochs — количество эпох. Начните с 300, но если модель переобучается (валидационная точность падает), уменьшите значение. Если переобучения нет, можно увеличить до 600 или 1200.
  • batch_size — размер партии. Установите максимальное значение, которое поддерживает ваша GPU. Если возникает ошибка нехватки памяти, уменьшайте его постепенно. В YOLO можно задать batch=-1, чтобы модель автоматически подобрала оптимальный размер.
  • optimizer — алгоритм оптимизации. По умолчанию используется SGD, но можно попробовать Adam или AdamW. Выбор оптимизатора влияет на скорость и стабильность обучения.
  • learning_rate — скорость обучения. Начните с умеренного значения (например, 0.01) и постепенно снижайте его в процессе обучения. Импульс (momentum) обычно устанавливают около 0.9.
  • patience — параметр ранней остановки. Если в течение указанного числа эпох (например, 5) метрики валидации не улучшаются, обучение прекращается. Это экономит время и предотвращает переобучение.
  • amp — смешанная точность (mixed precision). Установите amp=True, чтобы использовать FP16 и FP32 одновременно. Это ускоряет обучение и снижает использование памяти без потери точности.

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

Аугментация данных и предобученные веса

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

В YOLO аугментация включена по умолчанию, но вы можете настроить её параметры в конфигурации обучения. Например, можно задать вероятность поворота, величину сдвига и другие параметры.

Использование предобученных весов — ещё один важный приём. Вместо обучения модели с нуля вы начинаете с весов, которые уже обучены на большом наборе данных (например, COCO). Это называется трансферным обучением. Оно значительно сокращает время обучения и улучшает точность, так как модель уже понимает базовые признаки. В YOLO для этого достаточно указать путь к предобученным весам при загрузке модели, например YOLO('yolov8n.pt').

Обучение на GPU и в облаке

Обучение YOLO на больших датасетах требует значительных вычислительных ресурсов. Локальный GPU (например, NVIDIA RTX) — хороший вариант для небольших экспериментов, но для серьёзных проектов часто используют облачные платформы.

Облачные сервисы, такие как Google Cloud, AWS и Azure, предоставляют доступ к высокопроизводительным GPU и TPU по требованию. Это позволяет масштабировать обучение и проводить эксперименты с большими моделями. Однако облачное обучение может быть дорогим, особенно при длительном использовании, а передача данных увеличивает расходы и задержки.

При работе с GPU важно правильно настроить размер партии. Установите batch_size на максимальное значение, которое поддерживает ваша видеокарта, чтобы полностью использовать её возможности. Если возникают ошибки нехватки памяти, уменьшайте размер партии до тех пор, пока обучение не пойдёт стабильно.

Для локального обучения на слабом оборудовании можно использовать CPU, но это займёт значительно больше времени. В таком случае рекомендуется использовать предобученные веса и небольшие модели (например, yolov8n).

Оценка и валидация модели

После обучения важно оценить качество модели. Основные метрики для детекции объектов:

  • mAP (mean Average Precision) — средняя точность по всем классам. Чем выше, тем лучше.
  • Precision — доля верно предсказанных объектов среди всех предсказаний.
  • Recall — доля найденных объектов среди всех реальных объектов.

В YOLO метрики автоматически вычисляются на валидационной выборке после каждой эпохи. Вы можете просмотреть графики в папке runs/detect/train, а также протестировать модель на новых изображениях.

Если модель показывает низкую точность, проверьте:

  • Достаточно ли данных в датасете?
  • Правильно ли размечены объекты?
  • Не переобучается ли модель (проверьте разницу между train и val метриками)?
  • Правильно ли подобраны гиперпараметры?

Также важно протестировать модель на реальных данных, которые отличаются от обучающих. Это поможет выявить проблемы с обобщением.

Практические примеры и ограничения

Рассмотрим несколько примеров использования YOLO:

  • Детекция объектов на фото — загрузите модель, передайте изображение и получите рамки с классами. Это базовая задача, с которой начинают большинство проектов.
  • Обработка видео — YOLO может обрабатывать каждый кадр видео, рисуя рамки и сохраняя результат. Это полезно для видеонаблюдения.
  • Сегментация и удаление фона — модель сегментации позволяет выделить объект (например, человека) и заменить фон. Это востребовано в фоторедакторах и видеомонтаже.
  • Трекинг объектов — позволяет отслеживать движение объектов в видеопотоке, что используется в системах безопасности и аналитике.

Однако у YOLO есть ограничения:

  • Точность на мелких объектах — YOLO может хуже детектировать очень маленькие объекты, особенно на больших изображениях.
  • Переобучение — при малом датасете модель может запомнить обучающие примеры и плохо работать на новых данных.
  • Требования к оборудованию — обучение больших моделей требует мощных GPU, что может быть дорого.

Несмотря на эти ограничения, YOLO остаётся одним из лучших инструментов для задач реального времени благодаря балансу скорости и точности.

Вопросы и ответы

Сколько изображений нужно для обучения YOLO?

Рекомендуется минимум 100–200 изображений на каждый класс. Чем больше разнообразия (ракурсы, освещение, фон), тем лучше модель обобщает. Для сложных задач может потребоваться несколько тысяч изображений.

Какой формат данных нужен для обучения YOLO?

YOLO использует текстовые файлы с координатами в нормализованном виде: class x_center y_center width height. Каждое изображение сопровождается файлом .txt с таким же именем. Также нужен файл data.yaml с путями к папкам и списком классов.

Что такое трансферное обучение и зачем оно нужно?

Трансферное обучение — это использование предобученных весов (например, на датасете COCO) для старта обучения на своих данных. Это ускоряет обучение и повышает точность, так как модель уже знает базовые признаки. В YOLO достаточно указать путь к предобученным весам при загрузке модели.

Как выбрать количество эпох для обучения?

Хорошая отправная точка — 300 эпох. Если модель переобучается (валидационная точность падает), уменьшите количество. Если переобучения нет, можно увеличить до 600 или 1200. Используйте параметр patience для ранней остановки, чтобы экономить время.

Можно ли обучить YOLO на CPU?

Да, можно, но это займёт значительно больше времени, особенно на больших датасетах. Для экспериментов с небольшими данными и моделями типа yolov8n CPU подойдёт. Для серьёзного обучения рекомендуется использовать GPU или облачные сервисы.

Как улучшить точность модели YOLO?

Увеличьте размер и разнообразие датасета, используйте аугментацию данных, настройте гиперпараметры (batch_size, learning_rate, optimizer), попробуйте более крупную версию модели (например, yolov8x вместо yolov8n), а также проверьте качество разметки.

Что делать, если модель переобучается?

Уменьшите количество эпох, увеличьте аугментацию данных, добавьте регуляризацию (например, weight decay), уменьшите размер модели или используйте предобученные веса. Также проверьте, достаточно ли данных в датасете.