跳到主要内容

微调 YOLO11（万圣节版）

· 阅读需 5 分钟
Aditya Sahu
Aditya Sahu
AI @ Synaptics

目标检测已获得显著进展，YOLO 系列是实时检测任务的主力军。每一代 YOLO 都带来更高的准确率和速度，适用于安防、自动驾驶等多种场景。在这篇博客中，你将学习如何在一个定制的南瓜数据集上微调最新的 YOLO11 模型，针对 Astra Machina 的 NPU 优化模型，并在 Astra Machina 上运行实时视频推理。

image

为什么要微调？​

为了让模型识别像“南瓜”这样的自定义标签，必须在新的定制数据集上对预训练模型进行微调。YOLO11 已在 COCO 数据集上预训练，COCO 数据集包含了80个常见目标类别，如人、动物和车辆。但 COCO 缺少专用类别，如南瓜，因此模型不能直接识别这些对象。微调让模型将从 COCO 学到的一般检测能力转移到识别新的类别（这里是南瓜），同时保留对空间特征和边界的理解。

Image

数据集是关键！​

没错，数据集决定了微调的效果。一个准备充分的数据集能让模型识别多种环境下的对象，从而应对真实应用中的各种变化。

像 Roboflow 这样的平台提供了大量数据集和模型，涵盖目标检测、分类、语义分割等多种任务。在本例中，你将使用专为万圣节准备的南瓜数据集。你可以从 Roboflow 下载 南瓜数据集。该数据集已为多款 YOLO 模型做了格式化；这里你将使用 YOLO11。以下是数据集的直接下载链接：南瓜数据集。

确保你的目录结构如下：

├── pumpkin-dataset/
│   ├── train/
│   ├── valid/
│   ├── test/
│   └── data.yaml

同时，确保 data.yaml 文件正确设置了数据集所需的信息，包括训练、验证和测试数据的路径。请检查文件中是否包含正确的数据路径，并指定 names 和 nc（类别数）的名称，如下：

nc: 1
names: ['Pumpkin']

了解 YOLO11 微调​

在微调 YOLO 时，你会从一个已在 COCO 数据集上学会检测一般目标的预训练模型开始。微调会逐步调整模型权重，使其学会南瓜的特征。你用 COCO 的预训练权重初始化模型，跳过初层的训练，因为这些层提取了边缘和纹理等通用低级特征。后面的层则会调整以识别南瓜特有的特征。由于 YOLO11 的结构已为目标检测做了优化，你只需修改输出层来识别新的类别。

因此，从 Ultralytics 加载一个预训练的 YOLO11 小型目标检测模型，并在南瓜数据集上训练它。这里，图像分辨率设置为 320x320；你也可以根据需要修改。

model = YOLO('yolo11s.pt')
model.train(data='/content/pumpkin-dataset/data.yaml', epochs=100, imgsz=(320,320), name='pumpkin_yolo', augment=True )

主要参数：

  • data: 指向数据配置文件（data.yaml）的路径。
  • epochs: 训练的轮数（可根据需求调整）。
  • imgsz: 图像尺寸（默认是 320x320）。
  • name: 训练运行的名称。

评估与测试模型​

训练结束后，在测试集上评估模型。关键指标包括准确率、召回率和平均准确率（mAP）等，用于检测南瓜。数据集较小时，模型容易过拟合，因此监控这些指标有助于调整参数以达到最佳效果。训练完成后，你可以在验证集或测试集上评估模型。

image

针对 Astra Machina 优化模型​

当模型训练完毕后，下一步是保存模型并导出为 TensorFlow Lite (.tflite) 格式，作为 float32 模型。接着，使用 SyNAP 工具包针对 Astra Machina 的 NPU 优化模型。该工具会处理优化和量化，将模型转换为 .synap 格式，使其在 NPU 上高效运行。

在优化过程中，建议在 meta.yaml 文件中提供一个与模型分辨率匹配的数据集。例如，由于模型使用 320x320 的图像，确保数据集中包含该尺寸的图像。你可以获取一张 320x320 的图像，或下载并调整 bus.jpg 的大小。

同时，请确保模型分辨率与 meta.yaml 中 shape 和 format 的设置一致。

!synap_convert --model /content/best_float32.tflite --target SL1680 --meta /content/synap.yaml --out-dir yolo

在 Astra Machina 上运行推理​

当模型优化并编译为适用于 Astra Machina 的 NPU 后，就可以运行起来了。首先，将生成的 .synap 模型传输到 Machina 板上。

由于你专门训练模型去检测南瓜（只有一个类别），需要创建一个 label.json 文件来告诉系统该类别。文件内容很简单，只需将标签映射到索引，如下：

{
    "labels": ["Pumpkin"]
}

SyNAP Examples GitHub 库提供了配好的 Python 示例，支持实时视频推理。无论你使用实时摄像头、视频文件，还是 RTSP 流，这些脚本都能处理所有内容，让你专心看模型如何实时检测南瓜。

进一步阅读​

微调 YOLO11 以检测南瓜，并针对 Astra Machina 的 NPU 进行优化，展示了目标检测模型如何适应特定任务。对于那些想深入了解的人，这里有一个详细的分步教程。