Обучение 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 предлагает простое решение без кода для всего цикла операций машинного обучения (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: ../datasets/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 (Segment Anything Model 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("runs/train/exp/results.csv")  # В этом же каталоге будет создан файл 'results.png'

5. Дальнейшие шаги#

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

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

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

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

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

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

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

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