Обучение YOLOv5 на пользовательских данных#

📚 В этом руководстве объясняется, как обучить собственный пользовательский датасет с помощью модели YOLOv5 🚀. Обучение пользовательских моделей — фундаментальный этап адаптации решений в области компьютерного зрения под конкретные реальные задачи, выходящие за рамки общего обнаружения объектов.

Перед началом#

Сначала убедись, что необходимое окружение настроено. Клонируй репозиторий YOLOv5 и установи необходимые зависимости из requirements.txt. Обязательно используй окружение Python>=3.8.0 с PyTorch>=1.8. Модели и датасеты автоматически скачиваются из последнего релиза YOLOv5, если они не найдены локально.

git clone https://github.com/ultralytics/yolov5 # Clone the repository
cd yolov5
pip install -r requirements.txt # Install dependencies

Обучение на пользовательских данных#

Разработка пользовательской модели обнаружения объектов — это итеративный процесс:

1. Сбор и организация изображений: Собери изображения, соответствующие конкретной задаче. Высококачественные и разнообразные данные имеют решающее значение. См. наше руководство по сбору и аннотированию данных.
2. Разметка объектов: Точно аннотируй интересующие объекты на изображениях.
3. Обучение модели: Используй размеченные данные, чтобы обучить модель YOLOv5. Примени перенос обучения, начав с предварительно обученных весов.
4. Развертывание и предсказание: Используй обученную модель для выполнения инференса на новых, ранее не встречавшихся данных.
5. Сбор пограничных случаев: Определи сценарии, в которых модель работает плохо (пограничные случаи), и добавь в датасет похожие данные, чтобы повысить устойчивость. Повтори цикл.

Ultralytics Platform предлагает оптимизированное решение без кода для всего этого цикла операций машинного обучения (MLOps), включая управление датасетами, обучение моделей и развертывание.

Перед началом обучения необходимо подготовить датасет.

1. Создание датасета#

Моделям YOLOv5 нужны размеченные данные, чтобы изучать визуальные характеристики классов объектов. Правильная организация датасета имеет ключевое значение.

1.1 Создание dataset.yaml#

Файл конфигурации датасета (например, coco128.yaml) описывает структуру датасета, названия классов и пути к каталогам с изображениями. COCO128 — небольшой пример датасета, содержащий первые 128 изображений из большого датасета COCO. Он полезен для быстрой проверки конвейера обучения и диагностики потенциальных проблем, таких как переобучение.

Структура файла dataset.yaml включает:

• path: корневой каталог, содержащий датасет.
• train, val, test: относительные пути от path к каталогам с изображениями или текстовым файлам со списками путей к изображениям для обучающей, валидационной и тестовой выборок.
• names: словарь, сопоставляющий индексы классов (начиная с 0) соответствующим названиям классов.

Можно задать path как абсолютный каталог (например, /home/user/datasets/coco128) или относительный путь, например ../datasets/coco128, при запуске обучения из корневого каталога репозитория YOLOv5.

Ниже приведена структура coco128.yaml (см. на GitHub):

# Dataset root directory relative to the yolov5 directory
path: coco128

# Train/val/test sets: specify directories, *.txt files, or lists
train: images/train2017 # 128 images for training
val: images/train2017 # 128 images for validation
test: # Optional path to test images

# Classes (example using 80 COCO classes)
names:
    0: person
    1: bicycle
    2: car
    # ... (remaining COCO classes)
    77: teddy bear
    78: hair drier
    79: toothbrush

1.2 Использование моделей для автоматической разметки#

Ручная разметка — распространенный подход, но она отнимает много времени. Базовые модели могут автоматизировать или частично автоматизировать аннотирование и ускорить создание датасета. Примеры моделей, которые могут помочь с созданием разметки:

• Google Gemini: Большие мультимодальные модели, такие как Gemini, обладают мощными возможностями понимания изображений. Им можно задать запрос на обнаружение и локализацию объектов на изображениях, чтобы получить ограничивающие рамки или описания, которые можно преобразовать в метки формата YOLO. Изучи возможности этой модели в прилагаемом учебном ноутбуке.
• SAM2 (Модель сегментации любых объектов 2): Базовые модели, ориентированные на сегментацию, такие как SAM2, могут с высокой точностью обнаруживать объекты и очерчивать их границы. Хотя они предназначены главным образом для сегментации, полученные маски часто можно преобразовать в аннотации ограничивающих рамок для задач обнаружения объектов.
• YOLOWorld: Эта модель поддерживает обнаружение объектов в открытом словаре. Ты можешь указать текстовые описания интересующих объектов, и YOLOWorld найдет их на изображениях без предварительного обучения на этих конкретных классах. Это можно использовать как отправную точку для создания начальных меток, которые затем можно уточнить.

Использование этих моделей позволяет добавить этап «предварительной разметки» и сократить объем ручной работы. Однако крайне важно проверить и уточнить автоматически созданные метки, чтобы обеспечить точность и согласованность, поскольку их качество напрямую влияет на производительность обученной модели YOLOv5. После создания (и возможного уточнения) меток убедись, что они соответствуют формату YOLO: один файл *.txt на изображение, где каждая строка представляет объект в виде class_index x_center y_center width height (нормализованные координаты, индекс класса с отсчетом от нуля). Если изображение не содержит интересующих объектов, соответствующий файл *.txt не нужен.

Требования к файлу *.txt формата YOLO точно определены:

• Одна строка на каждый объект ограничивающей рамки.
• Каждая строка должна содержать: class_index x_center y_center width height.
• Координаты необходимо нормализовать в диапазон от 0 до 1. Для этого раздели значения пикселей x_center и width на общую ширину изображения, а y_center и height — на общую высоту изображения.
• Индексы классов отсчитываются от нуля (то есть первый класс обозначается как 0, второй — как 1 и так далее).

Файл меток, соответствующий изображению выше и содержащий два объекта «person» (индекс класса 0) и один объект «tie» (индекс класса 27), будет выглядеть так:

1.3 Организация каталогов#

Организуй каталог датасетов в соответствии с приведенной ниже схемой. По умолчанию YOLOv5 ожидает, что каталог датасета (например, /coco128) будет находиться в папке /datasets, расположенной рядом с каталогом репозитория /yolov5.

YOLOv5 автоматически находит метки для каждого изображения, заменяя последнее вхождение /images/ в пути к изображению на /labels/. Например:

../datasets/coco128/images/im0.jpg # Path to the image file
../datasets/coco128/labels/im0.txt # Path to the corresponding label file

Рекомендуемая структура каталогов:

/datasets/
└── coco128/  # Dataset root
    ├── images/
    │   ├── train2017/  # Training images
    │   │   ├── 000000000009.jpg
    │   │   └── ...
    │   └── val2017/    # Validation images (optional if using same set for train/val)
    │       └── ...
    └── labels/
        ├── train2017/  # Training labels
        │   ├── 000000000009.txt
        │   └── ...
        └── val2017/    # Validation labels (optional if using same set for train/val)
            └── ...

2. Выбор модели#

Выбери предварительно обученную модель, чтобы начать процесс обучения. Использование предварительно обученных весов значительно ускоряет обучение и повышает производительность по сравнению с обучением с нуля. YOLOv5 предлагает модели разных размеров, каждая из которых обеспечивает свой баланс скорости и точности. Например, YOLOv5s — вторая по компактности и самая быстрая модель, подходящая для окружений с ограниченными ресурсами. Подробное сравнение всех доступных моделей см. в таблице README.

3. Обучение#

Начни обучение модели с помощью скрипта train.py. Основные аргументы:

• --img: задает размер изображения на входе (например, --img 640). Больший размер обычно обеспечивает более высокую точность, но требует больше памяти GPU.
• --batch: определяет размер пакета (например, --batch 16). Выбери максимальный размер, который может обработать твой GPU.
• --epochs: задает общее количество обучающих эпох (например, --epochs 100). Одна эпоха представляет собой полный проход по всему обучающему датасету.
• --data: путь к файлу dataset.yaml (например, --data coco128.yaml).
• --weights: путь к файлу исходных весов. Настоятельно рекомендуется использовать предварительно обученные веса (например, --weights yolov5s.pt) для более быстрой сходимости и лучших результатов. Для обучения с нуля (не рекомендуется, если только у тебя нет очень большого датасета и особых требований) используй --weights '' --cfg yolov5s.yaml.

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

# Example: Train YOLOv5s on the COCO128 dataset for 3 epochs
python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt

Все результаты обучения, включая веса и логи, сохраняются в каталоге runs/train/. Каждый сеанс обучения создает новый подкаталог (например, runs/train/exp, runs/train/exp2 и т. д.). Для интерактивного практического опыта изучи раздел об обучении в официальных учебных ноутбуках:

4. Визуализация#

YOLOv5 легко интегрируется с различными инструментами для визуализации хода обучения, оценки результатов и мониторинга производительности в реальном времени.

Логирование и визуализация в Comet#

Comet полностью интегрирован для комплексного отслеживания экспериментов. Визуализируй метрики в реальном времени, сохраняй гиперпараметры, управляй датасетами и контрольными точками моделей, а также анализируй предсказания моделей с помощью интерактивных пользовательских панелей Comet.

Начать работу легко:

pip install comet_ml                                                          # 1. Install Comet library
export COMET_API_KEY=YOUR_API_KEY_HERE                                        # 2. Set your Comet API key (create a free account at Comet.ml)
python train.py --img 640 --epochs 3 --data coco128.yaml --weights yolov5s.pt # 3. Train your model - Comet automatically logs everything!

Подробнее о поддерживаемых функциях можно узнать в нашем руководстве по интеграции с Comet. Дополнительные сведения о возможностях Comet см. в официальной документации. Попробуй ноутбук Comet в Colab для демонстрации в реальном времени:

Логирование и автоматизация в ClearML#

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

• Установи пакет: pip install clearml
• Инициализируй ClearML: один раз выполни clearml-init, чтобы подключиться к серверу ClearML (самостоятельно размещенному или доступному на бесплатном тарифе).

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

Локальное логирование#

Результаты обучения автоматически записываются с помощью TensorBoard и сохраняются в виде файлов CSV в каталоге конкретного эксперимента (например, runs/train/exp). Записываемые данные включают:

• Потери и метрики производительности на этапах обучения и валидации.
• Примеры изображений с примененными аугментациями (например, мозаичными).
• Эталонные метки и предсказания модели для визуальной проверки.
• Ключевые метрики оценки, такие как кривые точности-полноты (PR).
• Матрицы ошибок для подробного анализа производительности по классам.

Файл results.csv обновляется после каждой эпохи и по завершении обучения отображается как results.png. Ты также можешь вручную построить график любого файла results.csv с помощью предоставленной служебной функции:

from utils.plots import plot_results

# Plot results from a specific training run directory
plot_results("runs/train/exp/results.csv")  # This will generate 'results.png' in the same directory

5. Следующие шаги#

После успешного завершения обучения контрольная точка лучшей модели (best.pt) сохраняется и готова к развертыванию или дальнейшему улучшению. Возможные следующие шаги:

• Выполни инференс на новых изображениях или видео с помощью обученной модели через CLI или Python.
• Выполни валидацию, чтобы оценить точность модели и ее способность к обобщению на разных частях данных (например, на отложенной тестовой выборке).
• Экспортируй модель в различные форматы для развертывания, например ONNX, TensorFlow SavedModel или TensorRT, чтобы оптимизировать инференс на разных платформах.
• Применяй методы настройки гиперпараметров, чтобы потенциально добиться дополнительного прироста производительности.
• Продолжай улучшать модель, следуя нашим советам для достижения наилучших результатов обучения и итеративно добавляя более разнообразные и сложные данные на основе анализа производительности.

Поддерживаемые среды#

Ultralytics предоставляет готовые окружения с необходимыми зависимостями, такими как CUDA, cuDNN, Python и PyTorch, что упрощает начало работы.

• Бесплатные ноутбуки с GPU:
  • 
  • 
  • 
• Облачные платформы:
  • Google Cloud: краткое руководство по GCP
  • Amazon AWS: краткое руководство по AWS
  • Microsoft Azure: краткое руководство по AzureML
• Локальная настройка:
  • Docker: Краткое руководство по Docker

Статус проекта#

Этот значок указывает, что все тесты GitHub Actions в рамках непрерывной интеграции (CI) YOLOv5 успешно проходят. Эти строгие тесты CI охватывают основные функции, включая обучение, валидацию, инференс, экспорт и бенчмарки, в операционных системах macOS, Windows и Ubuntu. Тесты выполняются автоматически каждые 24 часа и после каждого коммита кода, обеспечивая стабильность и оптимальную производительность.

Часто задаваемые вопросы#