YOLO-NAS:神经架构搜索的下一代 YOLO
YOLO-NAS 目标检测实战教程
YOLO-NAS(You Only Look Once – Neural Architecture Search)是 Deci AI 推出的新一代实时目标检测模型。它通过自动化神经架构搜索引擎 AutoNAC 找到精准度与速度之间的最优平衡,性能超越了此前主流的 YOLOv5、YOLOv6、YOLOv7 和 YOLOv8。
本教程面向零基础读者,将带你从概念理解到代码实战,完成图像和视频中的目标检测。
什么是 YOLO-NAS
YOLO-NAS 属于单阶段目标检测器家族,它在设计上继承了 YOLO 系列“只看一次”的高效推理,并引入了以下关键创新:
- 自动化架构搜索:利用 AutoNAC 自动探索数亿种可能的网络结构,找出在给定硬件和推理约束下最优的模型。
- 量化友好型设计:原生支持 INT8 量化,几乎无损地部署在边缘设备上。
- 后处理优化:内置改进的非极大值抑制(NMS)算法,提升密集场景下的检测精度。
YOLO-NAS 提供三种尺寸:yolo_nas_s(小型)、yolo_nas_m(中型)、yolo_nas_l(大型),分别对应不同的算力需求与精度。
环境搭建与安装
本教程基于 Python 和 PyTorch 环境。推荐使用虚拟环境以避免依赖冲突。
-
创建并激活虚拟环境(可选)
python -m venv yolo_nas_env source yolo_nas_env/bin/activate # Windows: yolo_nas_env\Scripts\activate -
安装 PyTorch
根据你的 CUDA 版本,从 PyTorch 官网 获取安装命令。例如,CUDA 11.8 版本:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 -
安装 SuperGradients
YOLO-NAS 的训练和推理通过 Deci 的 SuperGradients 库实现:
pip install super-gradients -
验证安装
import super_gradients print(super_gradients.__version__)若正常输出版本号,即安装成功。
使用预训练模型进行推理
SuperGradients 提供了在 COCO 数据集上预训练的 YOLO-NAS 模型,可直接调用。以下是最小示例。
模型加载与单张图像检测
from super_gradients.training import models
from super_gradients.common.object_names import Models
# 加载预训练模型(这里以 yolo_nas_s 为例)
model = models.get(Models.YOLO_NAS_S, pretrained_weights="coco")
# 执行推理,返回检测结果列表
results = model.predict("path/to/your/image.jpg")
# 可视化结果并保存
results.show() # 在窗口展示
results.save("output.jpg") # 保存到文件
results 对象包含了边界框坐标、类别 ID 和置信度分数。可以通过以下方式获取结构化数据:
predictions = results.prediction # 包含 bboxes_xyxy, confidence, labels 等
for img_pred in predictions:
print(img_pred.bboxes_xyxy)
print(img_pred.confidence)
print(img_pred.labels)
视频和批量推理
# 对视频文件进行检测
results = model.predict("video.mp4")
results.save("output_video.mp4") # 自动写入带检测框的视频
# 对目录下所有图像批量推理
import glob
image_list = glob.glob("images/*.jpg")
results = model.predict(image_list)
results.save(output_folder="output_images/")
常用预测参数
预测时可以传入额外参数控制行为:
conf:置信度阈值,默认0.25。iou:NMS 的 IoU 阈值,默认0.7。image_size:推理时图像尺寸,默认基于模型配置。
results = model.predict("image.jpg", conf=0.5, iou=0.6, image_size=640)
自定义数据集训练
如果需要在你的特定任务上训练 YOLO-NAS,需要准备数据并将其转换为 YOLO 格式或使用 SuperGradients 的数据加载器。
数据集格式
SuperGradients 支持多种标注格式,推荐使用 YOLO 格式:
- 目录结构:
dataset/ train/ images/ labels/ val/ images/ labels/ - 标签文件:每个
.txt对应一张图像,每行一个目标:class_id x_center y_center width height(归一化坐标)。
你也可以使用 COCO JSON 格式,训练时将 dataset_params 设置为 coco_detection_yolo_format。
配置训练参数
创建一个 YAML 或 Python 字典来定义训练配方(recipe)。SuperGradients 提供了内置的 YOLO-NAS 训练配方。
from super_gradients.training import Trainer, dataloaders
from super_gradients.training.dataloaders.dataloaders import coco_detection_yolo_format_val, coco_detection_yolo_format_train
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
# 定义数据集路径
dataset_params = {
'data_dir': 'path/to/dataset',
'train_images_dir': 'train/images',
'train_labels_dir': 'train/labels',
'val_images_dir': 'val/images',
'val_labels_dir': 'val/labels',
'classes': ['class1', 'class2', ...] # 你的类别列表
}
# 创建训练和验证数据加载器
train_loader = 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': 8,
'num_workers': 4
}
)
val_loader = 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': 8,
'num_workers': 4
}
)
# 加载模型并调整类别数
model = models.get(Models.YOLO_NAS_S, pretrained_weights="coco", num_classes=len(dataset_params['classes']))
# 定义训练参数
train_params = {
'max_epochs': 20,
'lr_mode': 'cosine',
'initial_lr': 0.001,
'loss': PPYoloELoss(
use_static_assigner=False,
num_classes=len(dataset_params['classes'])
),
'optimizer': 'AdamW',
'average_best_models': True,
'valid_metrics_list': [
DetectionMetrics_050(score_thres=0.1, top_k_predictions=300, post_prediction_callback=PPYoloEPostPredictionCallback()),
DetectionMetrics_050_095(score_thres=0.1, top_k_predictions=300, post_prediction_callback=PPYoloEPostPredictionCallback())
],
'metric_to_watch': 'mAP@0.50:0.95'
}
# 启动训练
trainer = Trainer(experiment_name='yolo_nas_custom', ckpt_root_dir='checkpoints')
trainer.train(model=model, training_params=train_params, train_loader=train_loader, valid_loader=val_loader)
训练完成后,最佳模型权重会保存在 checkpoints/yolo_nas_custom 目录下。
使用训练后的模型推理
best_model = models.get(Models.YOLO_NAS_S,
num_classes=len(dataset_params['classes']),
checkpoint_path="checkpoints/yolo_nas_custom/average_model.pth")
results = best_model.predict("test.jpg")
results.show()
模型导出与优化部署
YOLO-NAS 支持导出为 ONNX 和 TensorRT,以在无 Python 环境中高效运行。
导出 ONNX
model.export("yolo_nas_custom.onnx",
input_image_shape=(1, 3, 640, 640),
output_path="onnx_model")
使用 INT8 量化
SuperGradients 提供了量化工具:
from super_gradients.common.quantization import QuantizationManager
manager = QuantizationManager()
quantized_model = manager.quantize(model, calibration_data=train_loader, dataloader_bs=8)
results = quantized_model.predict("test.jpg")
量化后的模型推理速度显著提升,尤其适合边缘设备部署。
常见问题与技巧
- 内存不足:减小
batch_size或使用更小的模型(如yolo_nas_s)。 - 自定义类别映射:确保
classes列表顺序与标签文件中的class_id一致。 - 训练精度不收敛:尝试更长的训练轮次、调整学习率或使用预训练权重进行更稳的初始化。
- 推理置信度阈值:根据场景调整
conf,高阈值过滤低质量检测,低阈值保留更多候选框。
通过本教程,你已掌握 YOLO-NAS 的基础使用、自定义训练与部署流程。基于神经架构搜索的设计让它在无需手动调参的情况下就能获得先进的检测性能,是实时视觉应用的强力选择。