На данный момент YOLO-NAS является новейшей моделью обнаружения объектов YOLO. С самого начала он превосходит все другие модели YOLO по точности. Предварительно обученные модели YOLO-NAS обнаруживают больше объектов с большей точностью по сравнению с предыдущими моделями YOLO. Но как нам обучить YOLO NAS пользовательскому набору данных? Это и будет нашей целью в этой статье – обучить различные модели YOLO NAS пользовательскому набору данных.

Основная претензия YOLO-NAS заключается в том, что он может обнаруживать объекты меньшего размера лучше, чем предыдущие модели. Хотя мы можем провести несколько экспериментов с выводом для анализа результатов, обучение на сложном наборе данных даст нам лучшее понимание. С этой целью мы проведем четыре обучающих эксперимента с использованием трех доступных предварительно обученных моделей YOLO-NAS. Для этой цели мы выбираем набор данных для тепловизионного обнаружения БПЛА.

В процессе экспериментов мы пройдем полный конвейер обучения для YOLO-NAS.

Набор данных обнаружения объектов для обучения YOLO NAS

Давайте начнем с ознакомления с набором данных о высотном инфракрасном тепловом режиме БПЛА.

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

Набор данных содержит 2898 тепловизионных изображений по 5 классам объектов:

• Person

• Car

• Bicycle

• OtherVehicle

• DontCare

Набор данных уже содержит обучающие, проверочные и тестовые разбиения. Имеется 2008 обучающих, 287 проверочных и 571 тестовый образцы. Набор данных уже присутствует в формате аннотаций YOLO.

Ниже приведены некоторые из неназванных изображений ground truth из набора данных.

Совершенно очевидно, что, кроме автомобилей, человеческий глаз не может видеть другие объекты на земле без надлежащих примечаний.

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

Приведенный выше рисунок помогает лучше понять объекты набора данных. Это также показывает нам, насколько трудно человеческому глазу распознавать объекты, такие как люди и велосипеды, в этом наборе данных. Будем надеяться, что модель YOLO-NAS после обучения будет работать намного лучше, чем мы, люди.

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

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

• YOLO_NAS_Fine_Tuning.ipynb

• YOLO_NAS_Large_Fine_Tuning.ipynb

• inference.ipynb

Мы подробно рассмотрим YOLO_NAS_Fine_Tuning.ipynbноутбук. Эти два шага содержат все шаги, необходимые для обучения YOLO NAS пользовательскому набору данных и последующего выполнения вывода с использованием обученных моделей. Учебные тетради содержат код для загрузки набора данных.

Следующий код обучит три модели YOLO NAS:

• YOLO NAS s (маленький)

• YOLO NAS m (средний)

• YOLO NAS l (большой)

Перед началом работы вы можете установить super-gradients пакет, который нам понадобится в процессе обучения и вывода. Хотя записные книжки содержат команду для этого, вы также можете установить ее с помощью следующей команды:

pip install super-gradients

Точная настройка моделей YOLO NAS

Давайте углубимся в код notebook. Самый первый шаг - импортировать все необходимые пакеты и модули.

from super_gradients.training import Trainer
from super_gradients.training import dataloaders
from super_gradients.training.dataloaders.dataloaders import (
    coco_detection_yolo_format_train, 
    coco_detection_yolo_format_val
)
from super_gradients.training import models
from super_gradients.training.losses import PPYoloELoss
from super_gradients.training.metrics import (
    DetectionMetrics_050,
    DetectionMetrics_050_095
)
from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
from tqdm.auto import tqdm
 
import os
import requests
import zipfile
import cv2
import matplotlib.pyplot as plt
import glob
import numpy as np
import random

Давайте рассмотрим все основные импорта из приведенного выше блока кода.

• Trainer: Тренер инициирует процесс обучения и также настраивает каталог экспериментов.

• dataloaders: Super Gradients предлагает собственные загрузчики данных для беспроблемного обучения моделей YOLO NAS.

• coco_detection_yolo_format_traincoco_detection_yolo_format_val: Эти два помогут нам определить набор данных для обучения и проверки.

• models: Эта функция инициализирует модели YOLO NAS.

• PPYoloELoss: YOLO NAS использует PPYoloELoss во время обучения.

• DetectionMetrics_050DetectionMetrics_050_095: Во время обучения мы будем отслеживать карту на уровне 50% IoU и основной показатель.

Загрузка набора данных и структура каталогов

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

hit-uav
├── dataset.yaml
├── images
│   ├── test
│   ├── train
│   └── val
└── labels
    ├── test
    ├── train
    └── val

Параметры набора данных для обучения YOLO NAS

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

ROOT_DIR = 'hit-uav'
train_imgs_dir = 'images/train'
train_labels_dir = 'labels/train'
val_imgs_dir = 'images/val'
val_labels_dir = 'labels/val'
test_imgs_dir = 'images/test'
test_labels_dir = 'labels/test'
classes = ['Person', 'Car', 'Bicycle', 'OtherVechicle', 'DontCare']
 
dataset_params = {
    'data_dir':ROOT_DIR,
    'train_images_dir':train_imgs_dir,
    'train_labels_dir':train_labels_dir,
    'val_images_dir':val_imgs_dir,
    'val_labels_dir':val_labels_dir,
    'test_images_dir':test_imgs_dir,
    'test_labels_dir':test_labels_dir,
    'classes':classes 
}

Используя dataset_params, мы можем позже легко создавать наборы данных в требуемом формате. Кроме того, нам нужно определить количество периодов обучения, размер пакета и количество рабочих для обработки данных.

# Global parameters.
EPOCHS = 50
BATCH_SIZE = 16
WORKERS = 8

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

Подготовка набора данных для обучения YOLO NAS

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

train_data = coco_detection_yolo_format_train(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['train_images_dir'],
        'labels_dir': dataset_params['train_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size':BATCH_SIZE,
        'num_workers':WORKERS
    }
)
 
val_data = coco_detection_yolo_format_val(
    dataset_params={
        'data_dir': dataset_params['data_dir'],
        'images_dir': dataset_params['val_images_dir'],
        'labels_dir': dataset_params['val_labels_dir'],
        'classes': dataset_params['classes']
    },
    dataloader_params={
        'batch_size':BATCH_SIZE,
        'num_workers':WORKERS
    }
)

В приведенном выше блоке кода мы создаем train_data и val_data для обучения и проверки, соответственно.

Как мы можем видеть, мы используем dataset_params словарь для получения пути к набору данных и значений класса набора данных. Мы также определяем dataloader_params для каждого из них, который принимает размер пакета и количество рабочих элементов для процесса загрузки данных.

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

Определение преобразований и дополнений для обучения YOLO NAS

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

Конвейер преобразования данных Super Gradients предлагает эффективный способ применения дополнений и управления ими.

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

train_data.dataset.transforms

Мы увидим следующий результат:

[DetectionMosaic('additional_samples_count': 3, 'non_empty_targets': False, 'prob': 1.0, 'input_dim': [640, 640], 'enable_mosaic': True, 'border_value': 114),
 
 DetectionRandomAffine('additional_samples_count': 0, 'non_empty_targets': False, 'degrees': 10.0, 
'translate': 0.1, 'scale': [0.1, 2], 'shear': 2.0, 'target_size': [640, 640], 'enable': True, 'filter_box_candidates': True, 'wh_thr': 2, 'ar_thr': 20, 'area_thr': 0.1, 'border_value': 114),
 
 DetectionMixup('additional_samples_count': 1, 'non_empty_targets': True, 'input_dim': [640, 640], 'mixup_scale': [0.5, 1.5], 'prob': 1.0, 'enable_mixup': True, 'flip_prob': 0.5, 'border_value': 114),
.
.
.
XYCoordinateFormat object at 0x7efcc5c07340>), ('labels', name=labels length=1)]), 'output_format': OrderedDict([('labels', name=labels length=1), ('bboxes', name=bboxes length=4 format=
<super_gradients.training.datasets.data_formats.bbox_formats.cxcywh.CXCYWHCoordinateFormat object at 0x7efcc5976170>)]), 'max_targets': 120, 'min_bbox_edge_size': 1, 'input_dim': [640, 640], 'targets_format_converter': 
<super_gradients.training.datasets.data_formats.format_converter.ConcatenatedTensorFormatConverter object at 0x7efcb956b5b0>)]

Среди множества дополнений одним из них является MixUp. Однако из-за путаницы набор данных может быть чрезвычайно сложным для восприятия в моделях глубокого обучения, если с ним не обращаться должным образом. В результате первоначальных экспериментов мы обнаружили, что увеличение путаницы для этого набора данных делает набор данных очень запутанным, поскольку он накладывает одно изображение на другое. Удаление дополнения Mixup и сохранение остальных без изменений, казалось, улучшило производительность.

Мы можем удалить дополнение MixUp из списка, удалив элемент из индекса 2.

############ An example on how to modify augmentations ###########
train_data.dataset.transforms.pop(2)

Мы также можем визуализировать преобразованные изображения, используя следующий код.

train_data.dataset.plot(plot_transformed_data=True)

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

Параметры обучения YOLO NAS

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

train_params = {
    'silent_mode': False,
    "average_best_models":True,
    "warmup_mode": "linear_epoch_step",
    "warmup_initial_lr": 1e-6,
    "lr_warmup_epochs": 3,
    "initial_lr": 5e-4,
    "lr_mode": "cosine",
    "cosine_final_lr_ratio": 0.1,
    "optimizer": "Adam",
    "optimizer_params": {"weight_decay": 0.0001},
    "zero_weight_decay_on_bias_and_bn": True,
    "ema": True,
    "ema_params": {"decay": 0.9, "decay_type": "threshold"},
    "max_epochs": EPOCHS,
    "mixed_precision": True,
    "loss": PPYoloELoss(
        use_static_assigner=False,
        num_classes=len(dataset_params['classes']),
        reg_max=16
    ),
    "valid_metrics_list": [
        DetectionMetrics_050(
            score_thres=0.1,
            top_k_predictions=300,
            num_cls=len(dataset_params['classes']),
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(
                score_threshold=0.01,
                nms_top_k=1000,
                max_predictions=300,
                nms_threshold=0.7
            )
        ),
        DetectionMetrics_050_095(
            score_thres=0.1,
            top_k_predictions=300,
            num_cls=len(dataset_params['classes']),
            normalize_targets=True,
            post_prediction_callback=PPYoloEPostPredictionCallback(
                score_threshold=0.01,
                nms_top_k=1000,
                max_predictions=300,
                nms_threshold=0.7
            )
        )
    ],
    "metric_to_watch": 'mAP@0.50:0.95'
}

Во время обучения на выходе будет отображаться карта как с 50%-ным, так и с 50-95%-ным IoU. Однако мы отслеживаем только первичную метрику (mAP@0.50 : 0,95 IoU). Таким образом, лучшая модель будет сохранена в соответствии с этим.

Обучение модели YOLO NAS

Поскольку мы обучаем три разные модели, нам нужно немного автоматизировать процесс. Мы можем определить список, содержащий три названия моделей, и настроить каталог контрольных точек в соответствии с этим. При этом также будет загружена соответствующая модель, поскольку названия моделей в списке совпадают с названиями из super-gradients API.

models_to_train = [
    'yolo_nas_s',
    'yolo_nas_m',
    'yolo_nas_l'
]
 
CHECKPOINT_DIR = 'checkpoints'
 
for model_to_train in models_to_train:
    trainer = Trainer(
        experiment_name=model_to_train, 
        ckpt_root_dir=CHECKPOINT_DIR
    )
 
    model = models.get(
        model_to_train, 
        num_classes=len(dataset_params['classes']), 
        pretrained_weights="coco"
    )
 
    trainer.train(
        model=model, 
        training_params=train_params, 
        train_loader=train_data, 
        valid_loader=val_data
    )

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

Анализ результатов точной настройки YOLO NAS

Во время процесса обучения выходная ячейка / терминал отображает полное представление процесса обучения.

SUMMARY OF EPOCH 0
├── Training
│   ├── Ppyoloeloss/loss = 3.8575
│   ├── Ppyoloeloss/loss_cls = 2.3712
│   ├── Ppyoloeloss/loss_dfl = 1.1773
│   └── Ppyoloeloss/loss_iou = 0.3591
└── Validation
    ├── F1@0.50 = 0.0
    ├── F1@0.50:0.95 = 0.0
    ├── Map@0.50 = 0.0012
    ├── Map@0.50:0.95 = 0.0005
    ├── Ppyoloeloss/loss = 3.7911
    ├── Ppyoloeloss/loss_cls = 2.5251
    ├── Ppyoloeloss/loss_dfl = 0.9791
    ├── Ppyoloeloss/loss_iou = 0.3106
    ├── Precision@0.50 = 0.0
    ├── Precision@0.50:0.95 = 0.0
    ├── Recall@0.50 = 0.0
    └── Recall@0.50:0.95 = 0.0
.
.
.
SUMMARY OF EPOCH 50
├── Training
│   ├── Ppyoloeloss/loss = 1.4382
│   │   ├── Best until now = 1.433  (↗ 0.0053)
│   │   └── Epoch N-1      = 1.433  (↗ 0.0053)
│   ├── Ppyoloeloss/loss_cls = 0.6696
│   │   ├── Best until now = 0.6651 (↗ 0.0046)
│   │   └── Epoch N-1      = 0.6651 (↗ 0.0046)
│   ├── Ppyoloeloss/loss_dfl = 0.6859
│   │   ├── Best until now = 0.6846 (↗ 0.0013)
│   │   └── Epoch N-1      = 0.686  (↘ -0.0)
│   └── Ppyoloeloss/loss_iou = 0.1703
│       ├── Best until now = 0.17   (↗ 0.0003)
│       └── Epoch N-1      = 0.17   (↗ 0.0003)
└── Validation
    ├── F1@0.50 = 0.292
    │   ├── Best until now = 0.3025 (↘ -0.0104)
    │   └── Epoch N-1      = 0.2774 (↗ 0.0146)
    ├── F1@0.50:0.95 = 0.1859
    │   ├── Best until now = 0.1928 (↘ -0.007)
    │   └── Epoch N-1      = 0.1761 (↗ 0.0097)
    ├── Map@0.50 = 0.7631
    │   ├── Best until now = 0.7745 (↘ -0.0114)
    │   └── Epoch N-1      = 0.7159 (↗ 0.0472)
    ├── Map@0.50:0.95 = 0.4411
    │   ├── Best until now = 0.4443 (↘ -0.0032)
    │   └── Epoch N-1      = 0.4146 (↗ 0.0265)
    ├── Ppyoloeloss/loss = 1.5389
    │   ├── Best until now = 1.5404 (↘ -0.0015)
    │   └── Epoch N-1      = 1.5526 (↘ -0.0137)
    ├── Ppyoloeloss/loss_cls = 0.6893
    │   ├── Best until now = 0.687  (↗ 0.0024)
    │   └── Epoch N-1      = 0.6972 (↘ -0.0079)
    ├── Ppyoloeloss/loss_dfl = 0.7148
    │   ├── Best until now = 0.7136 (↗ 0.0012)
    │   └── Epoch N-1      = 0.7234 (↘ -0.0086)
    ├── Ppyoloeloss/loss_iou = 0.1969
    │   ├── Best until now = 0.1953 (↗ 0.0016)
    │   └── Epoch N-1      = 0.1975 (↘ -0.0006)
    ├── Precision@0.50 = 0.1828
    │   ├── Best until now = 0.1926 (↘ -0.0097)
    │   └── Epoch N-1      = 0.1718 (↗ 0.011)
    ├── Precision@0.50:0.95 = 0.1166
    │   ├── Best until now = 0.1229 (↘ -0.0063)
    │   └── Epoch N-1      = 0.1092 (↗ 0.0074)
    ├── Recall@0.50 = 0.8159
    │   ├── Best until now = 0.8939 (↘ -0.0781)
    │   └── Epoch N-1      = 0.8307 (↘ -0.0149)
    └── Recall@0.50:0.95 = 0.522
        ├── Best until now = 0.5454 (↘ -0.0234)
        └── Epoch N-1      = 0.5236 (↘ -0.0016)
 
===========================================================

Мы можем просмотреть журналы Tensorboard и проверить графики mAP для упрощения сравнения всех трех процессов обучения.

Журналы Tensorboard присутствуют в соответствующих папках обучения внутри checkpoints каталога.

На следующем графике показано первичное сравнение точек доступа для трех обучающих экспериментов.

На приведенном выше графике:

• Красная строка: обучение большой модели YOLO NAS

• Синяя линия: обучение модели среды YOLO NAS

• Оранжевая линия: обучение модели YOLO NAS на малой основе

Большая модель YOLO NAS достигла самого высокого значения 44,4% mAP в на 43 эпохе. Здесь следует отметить еще один момент: большая модель YOLO NAS сравнительно раньше достигла более высокой карты. Это указывает на его способность обнаруживать сложные объекты по сравнению со средними и малыми моделями YOLO NAS.

super-gradients API сохраняет три разные контрольные точки для каждого эксперимента. Один для лучшей модели, один для последней модели и один для усредненных весов.

Поскольку большая модель YOLO NAS показала наилучшие результаты во время обучения пользовательскому набору данных, мы будем использовать ее в дальнейшем для экспериментов с выводом.

Тренируйте большую модель YOLO NAS дольше

Из приведенных выше экспериментов ясно, что большая модель YOLO NAS работает наилучшим образом. Чтобы получить еще лучшие результаты, мы можем обучать эту модель в течение 100 эпох. YOLO_NAS_Large_Fine_Tuning.ipynb Это достигается.

Вот график карты после обучения модели для 100 эпох на пользовательском наборе данных.

Точная настройка большой модели YOLO NAS в пользовательском наборе данных позволяет получить более 45% mAP.

Вывод на тестовых изображениях с использованием обученной модели YOLO NAS

Набор данных содержал тестовое разделение, которое мы сохранили зарезервированным для целей вывода. Вы можете выполнить ячейки кода в inference.ipynb записной книжке, чтобы запустить эксперимент вывода. Это позволяет выполнить несколько задач:

• Сначала он загружает наиболее подготовленные веса YOLO NAS из каталога контрольных точек.

• Затем он выполняет вывод на основе тестовых изображений. При этом код сохраняет результаты вывода в inference_results/images каталоге с оригинальными именами изображений.

• После получения результатов в записной книжке отображается набор изображений путем наложения аннотаций "Истинность земли" на прогнозируемые изображения.

На последнем шаге мы узнаем, какие объекты пропустила обученная модель и делала ли модель какие-либо ложные прогнозы.

Давайте начнем наш анализ с визуализации нескольких прогнозов вывода.

При самом начальном анализе кажется, что модель предсказывает почти все объекты, даже людей, которые кажутся очень маленькими. Но здесь чрезвычайно сложно указать на какие-либо ошибки.

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

На приведенных выше изображениях аннотации к основной правде выделены красным цветом с именами классов внизу.

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

В целом модель работает хорошо и не обнаруживает только очень трудно распознаваемые объекты.

Результаты видеовыхода обученной модели YOLO NAS

Мы также провели эксперимент с выводом на основе тепловизионного видео с беспилотника.

^{Клип 1. Видео-прогнозы после обучения модели YOLO NAS пользовательскому набору данных.}

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

Вывод видеосигнала выполнялся на ноутбуке с графическим процессором GTX 1060, и модель работала в среднем со скоростью 17 кадров в секунду. Учитывая, что здесь мы используем большую модель YOLO NAS, мы получаем приличную скорость.

Заключение

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

Из экспериментов мы можем сделать вывод, что новые модели YOLO NAS открывают новый мир для обнаружения в режиме реального времени. К ним относятся видеонаблюдение, мониторинг трафика и медицинская визуализация, среди многих других.

Код на GitHub