训练自己的YOLO-NAS

YOLO-NAS 是 Deci 开发的一种新的最先进的目标检测模型。 在本指南中，我们将讨论什么是 YOLO-NAS 以及如何在自定义数据集上训练 YOLO-NAS 模型。

为了训练我们的自定义模型，我们将：

• 加载预训练的YOLO-NAS模型；
• 从 Roboflow 加载自定义数据集，或者使用UnrealSynth制作合成数据集
• 设置超参数值；
• 使用超级梯度 Python 包根据我们的数据训练模型；
• 评估模型以了解结果。

话不多说，让我们开始吧！

1、什么是 YOLO-NAS？

You Only Look Once  神经架构搜索（YOLO-NAS）是最新最先进的（SOTA）实时目标检测模型。 在 COCO 数据集上进行评估并与其前身 YOLOv6 和 YOLOv8  相比，YOLO-NAS 以更低的延迟实现了更高的 mAP 值。

YOLO-NAS 作为 Deci 维护的 super-gradient包的一部分提供。

下图展示了Deci在YOLO-NAS上的基准测试结果：

YOLO-NAS 在 Roboflow 100 数据集基准测试中也是最好的，这表明它可以轻松地在自定义数据集上进行微调。

2、Python环境设置

在开始训练之前，我们需要准备好Python环境。 让我们从安装三个 pip 包开始。 YOLO-NAS 模型本身是使用 super-gradient 包进行分发的。 请记住，该模型仍在积极开发中。 为了保持环境的稳定性，最好固定特定版本的包。 此外，我们将安装 roboflow 和监督，这将使我们能够从 Roboflow Universe 下载数据集并分别可视化我们的训练结果。

pip install super-gradients==3.1.1
pip install roboflow
pip install supervision

如果你在 Jupyter Notebook 中运行 YOLO-NAS，请不要忘记在安装完成后重新启动环境。

3、使用预训练模型进行推理

在开始培训之前，最好确保安装按计划进行。 最简单的方法是使用预先训练的模型之一进行测试推理。 同时，这也能让我们熟悉YOLO-NAS API。

3.1 加载YOLO-NAS模型

为了使用预训练的 COCO 模型进行推理，我们首先需要选择模型的大小。 YOLO-NAS提供三种不同的模型大小：yolo_nas_s、yolo_nas_m和yolo_nas_l。

yolo_nas_s 模型是最小且最快的，但它可能不会像较大的模型那么准确。 相反，yolo_nas_l 模型最大、最准确、最慢。 yolo_nas_m 模型提供了两者之间的中间立场。

import torch
from super_gradients.training import models

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
MODEL_ARCH = 'yolo_nas_l'
#            'yolo_nas_m'
#            'yolo_nas_s'

model = models.get(MODEL_ARCH, pretrained_weights="coco").to(DEVICE)

3.2 YOLO-NAS模型推理

推理过程包括设置置信度阈值和调用预测方法。 预测方法将返回预测列表，其中每个预测对应于图像中检测到的对象。

CONFIDENCE_TRESHOLD = 0.35

result = list(model.predict(image, conf=CONFIDENCE_TRESHOLD))[0]

3.3 YOLO-NAS 推理输出格式

YOLO-NAS 推理的输出是一个 ImageDetectionPrediction 对象，它封装了图像中检测到的对象的详细信息。 该对象包含三个字段：

• image - 表示用于推理的图像的 NumPy 数组。
• class_names - 模型训练期间使用的类别名称的 Python 列表。
• Prediction -DetectionPrediction 类的实例，其中包含有关模型检测的详细信息。

DetectionPrediction对象具有三个字段：

• bboxes_xyxy - 形状 (N, 4) 的 NumPy 数组，以 xyxy 格式表示检测到的对象的边界框。
• confidence - 形状 (N,) 的 NumPy 数组，表示检测的置信度值。 每个值都在 0 和 1 之间。
• labels - 形状 (N,) 的 NumPy 数组，表示检测到的对象的类 ID。 每个类 ID 对应于 class_names 列表中的一个索引。

4、使用开源数据集微调 YOLO-NAS

为了微调模型，我们需要数据。 我们将使用足球运动员检测图像数据集。

如果你已经有 YOLO 格式的数据集，请随意使用它。 如果没有，请看看 Roboflow Universe，那里拥有超过 200,000 个开源项目，并且所有项目都可以以任何格式导出。

另外一种获取数据集的方法是使用UnrealSynth，一个基于虚幻引擎开发的YOLO合成数据生成器，可以自动生成包括标注的训练数据集，非常方便：

import roboflow
from roboflow import Roboflow

roboflow.login()

rf = Roboflow()
project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
dataset = project.version(PROJECT_VERSION).download("yolov5")

要训练 YOLO-NAS 模型，你需要设置几个关键参数。

首先，你需要选择模型尺寸。 有三个选项可供选择：小型、中型和大型。 请记住，较大的模型可能需要更长的时间来训练并需要更多的内存，因此如果使用的资源有限，你可能需要考虑使用较小的模型。

接下来，你需要设置批量大小。 该参数指示在训练过程的每次迭代期间将有多少图像通过神经网络。 较大的批量大小将加快训练过程，但也需要更多的内存。

MODEL_ARCH = 'yolo_nas_l'
BATCH_SIZE = 8
MAX_EPOCHS = 25
CHECKPOINT_DIR = f'{HOME}/checkpoints'
EXPERIMENT_NAME = project.name.lower().replace(" ", "_")
LOCATION = dataset.location
CLASSES = sorted(project.classes.keys())

dataset_params = {
    'data_dir': LOCATION,
    'train_images_dir':'train/images',
    'train_labels_dir':'train/labels',
    'val_images_dir':'valid/images',
    'val_labels_dir':'valid/labels',
    'test_images_dir':'test/images',
    'test_labels_dir':'test/labels',
    'classes': CLASSES
}

from super_gradients.training.dataloaders.dataloaders import (
    coco_detection_yolo_format_train, coco_detection_yolo_format_val)

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': 2
    }
)

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': 2
    }
)

最后，你需要设置训练过程的纪元数。 这本质上是整个数据集通过神经网络的次数。

5、训练自定义 YOLO-NAS 模型

你可能已经注意到，训练模型的过程比 YOLOv8 更加冗长。 Ultralytics 模型中的许多功能需要在 CLI 中传递参数，而对于 YOLO-NAS，则需要编写自定义逻辑。

最后，我们准备开始训练。 在调用 train 方法之前，值得运行 TensorBoard。 这将使我们能够实时跟踪培训的关键指标。 值得一提的是，YOLO-NAS还支持W&B等最流行的实验记录仪。

trainer.train(
    model=model, 
    training_params=train_params, 
    train_loader=train_data, 
    valid_loader=val_data
)

6、评估自定义 YOLO-NAS 模型

训练结束后，你可以使用Trainer提供的测试方法评估模型的性能。 你需要传入测试集数据加载器，训练器将返回一个指标列表，包括通常用于评估对象检测模型的平均精度（mAP）。

trainer.test(
    model=best_model,
    test_loader=test_data,
    test_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
        )
    )
)

此外，你可以对测试集图像进行推理并可视化结果，以更好地了解模型在各个示例上的表现。 你还可以计算混淆矩阵，以更详细地了解每个类别的模型性能：

7、结束语

一夜之间，YOLO-NAS 成为实时物体检测器的新选择。 在为你的项目微调模型时，请记住要考虑所有方面——从模型准确性到推理速度，再到易于训练和许可限制。

原文链接：How to Train YOLO-NAS on a Custom Dataset

BimAnt翻译整理，转载请标明出处