Обучение 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 (Segment Anything Model 2): Базовые модели для сегментации, такие как SAM2, могут с высокой точностью находить и выделять объекты. Хотя они предназначены в основном для сегментации, полученные маски часто можно преобразовать в аннотации в виде ограничивающих рамок, подходящие для задач обнаружения объектов.
• YOLOWorld: Эта модель поддерживает обнаружение с открытым словарем. Ты можешь предоставить текстовые описания интересующих тебя объектов, и YOLOWorld сможет найти их на изображениях без предварительного обучения на этих конкретных классах. Это можно использовать как отправную точку для генерации начальных меток, которые затем можно доработать.

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

Спецификации файлов разметки YOLO *.txt точны:

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

Файл разметки, соответствующий изображению выше и содержащий два объекта «человек» (индекс класса 0) и один объект «галстук» (индекс класса 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). Залогированные данные включают:

• Функции потерь при обучении и валидации, а также метрики производительности.
• Примеры изображений с примененными аугментациями (такими как мозаика).
• Метки истинных значений (Ground truth) вместе с предсказаниями модели для визуальной проверки.
• Ключевые метрики оценки, такие как кривые точности-полноты (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

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

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

FAQ#