Обучение модели детекции YOLO на своих данных

Ссылка на блокнот в Google Colab: https://colab.research.google.com/drive/1OUZ4wyjJj5xNzzUAB8hXY_Sl6uOfKNYg?usp=sharing

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

Для обучения критичную роль имеет качество и объем датасета, в обычных условиях выполняется правило: чем больше изображений, тем выше качество обучения. Минимальное количество картинок для использования данного скрипта, на мой взгляд, начинается от 100.

Сначала устанавливаем модуль «ultralytics»:

# Устанавливаем ultralytics
!pip install ultralytics

Затем подключаем необходимые модули и также даем блокноту доступ к Google Диску:

# Подключаем необходимые модули
from ultralytics import YOLO
import os
# Подключаем Google Диск
from google.colab import drive
drive.mount('/content/drive')

Далее указываем основные пути к данным для обучения:

# Устанавливаем папку с данными обучения
current_dir = os.path.dirname('/content/drive/MyDrive/ml/detection/yolo/train2/')
# Путь к конфигурационному файлу
data_path = os.path.join(current_dir, 'data.yaml')
# Путь к модели, на базе которой проводится обучение
model = YOLO(os.path.join(current_dir, 'yolov8n.pt'))

Здесь мы сначала указываем путь к директории, в которой находятся данные, в примере это «/ml/detection/yolo/train2/» (путь от корня Google Диска), структура данной папки следующая:

- yolov8n.pt
- data.yaml
- dataset
 - train
  - images
  - labels
 - test
  - images
  - labels
 - valid
  - images
  - labels
- output

Файл yolov8n.pt — это модель, используемая в качестве базовой для обучения, ее можно предварительно скачать (например, на https://huggingface.co/Ultralytics/YOLOv8/blob/main/yolov8n.pt) и загрузить в рабочую диекторию.

Файл data.yaml — это конфигурация для настройки обучения, его содержимое примерно следующее:

train: dataset/train/images
val: dataset/valid/images
test: dataset/test/images
nc: 1
names: [my_class]

Здесь указываются пути к изображениям, а также параметр «nc» — количество используемых классов (для новых объектов, находить которые учим) и названия этих классов (в примере класс 1 — то есть на всех новых изображениях указываем один объект, который хотим модель научить находить).

В папку «dataset» — добавляем три директории: «train», «valid», «test» и распределяем изображения по ним в процентной пропорции от 60/20/20 до 80/10/10 (можно экспериментировать с этим параметром). Каждая директория содержит 2 поддиректории: «images» — где лежат собственно сами изображения, «labels» — где лежат txt файлы с разметкой объектов (названия картинок и txt файлов должны совпадать, не считая расширение, чтобы система смогла сопоставить файлы, пример: «dataset/train/images/1.jpg» и «dataset/train/labels/1.txt»).

Файлы с разметкой имеют примерно следующее содержание:

0 0.6265625 0.50546875 1.253125 0.5140625

Где вначале идет порядковый номер класса объекта (в нашем 0 — так как объект единственный) и далее его координаты на изображении. Создать такие файлы можно вручную, либо воспользовавшись любым доступным инструментом, более подробно можно узнать по ссылке: https://docs.ultralytics.com/ru/yolov5/tutorials/train_custom_data/#11-create-datasetyaml

Далее в нашем коде устанавливаем параметры обучения:

# Параметры обучения
epochs = 250 # Количество эпох (циклов обучения)
batch = 64 # Размер пакета для обучения
imgsz = 640 # Размер изображений для обучения
patience = 0 # Параметр, который отвечает за остановку обучения при выходе на плато

Здесь значения «epochs» и «batch» — необходимо подбирать исходя из текущего объема датасета, а также доступной вычислительной мощности.

Параметр «patience» позволяет завершить обучение досрочно, если на протяжении нескольких циклов не будет улучшения, на небольших датасетах я предпочитаю его отключать (значение 0).

Завершает скрипт запуск обучения:

# Запуск обучения
if __name__ == '__main__':
    results = model.train(data=data_path,
                      epochs=epochs,
                      patience=patience,
                      batch=batch,
                      imgsz=imgsz,
                      name='/content/drive/MyDrive/ml/detection/yolo/train2/output/',
                      device='cuda')

После завершения в директории «/ml/detection/yolo/train2/output/» (от корня Google диска) должны создаться файлы с результатами обучения, в подпапке «weights» — файлы обученных моделей («best.pt» — модель с наилучшими показателями за обучение).