当前位置：首页 > news >正文

从入门到精通：EdgeAI-YOLOv5预训练模型使用与自定义数据集训练

news 2026/7/29 21:06:58

从入门到精通：EdgeAI-YOLOv5预训练模型使用与自定义数据集训练

【免费下载链接】edgeai-yolov5This repository has been deprecated. Suggest to use edgeai-mmdetection in https://github.com/TexasInstruments/edgeai-tensorlab项目地址: https://gitcode.com/gh_mirrors/ed/edgeai-yolov5

EdgeAI-YOLOv5是一款高效的目标检测工具，结合了YOLOv5的强大性能与边缘计算优化，特别适合在资源受限的设备上部署。本文将详细介绍如何使用预训练模型进行快速检测，以及如何训练自定义数据集，帮助新手轻松掌握目标检测技术。

🌟 EdgeAI-YOLOv5核心优势解析

EdgeAI-YOLOv5在保持高精度的同时，针对边缘设备进行了深度优化。通过对比YOLOv3、官方YOLOv5和YOLOv5-ti-lite的网络结构，可以清晰看到其在计算效率上的提升：

图1：YOLOv3、官方YOLOv5与YOLOv5-ti-lite网络结构对比，展示了EdgeAI版本在特征提取和计算效率上的优化设计

在性能表现上，EdgeAI-YOLOv5系列模型展现了优异的精度-速度平衡。从mAP（平均精度均值）与FLOPS（浮点运算次数）的关系图可以看出，即使在低计算资源下，模型仍能保持较高的检测精度：

图2：不同配置的EdgeAI-YOLOv5模型在COCO数据集上的mAP与FLOPS关系，绿色圆点表示性能最优的平衡配置

🚀 快速开始：预训练模型使用指南

1️⃣ 环境准备

首先克隆项目仓库并安装依赖：

git clone https://gitcode.com/gh_mirrors/ed/edgeai-yolov5 cd edgeai-yolov5 pip install -r requirements.txt

2️⃣ 下载预训练模型

项目提供了多种预训练模型，位于pretrained_models/models/detection/coco/edgeai-yolov5/目录下，包括：

yolov5s6_640_ti_lite_37p4_56p0.pt（轻量级模型，适合边缘设备）
yolov5m6_640_ti_lite_44p1_62p9.pt（中等规模模型，平衡精度与速度）
yolov5l6_640_ti_lite_47p1_65p6.pt（大型模型，高精度检测）

3️⃣ 运行目标检测

使用预训练模型对示例图片进行检测：

python detect.py --weights pretrained_models/models/detection/coco/edgeai-yolov5/yolov5s6_640_ti_lite_37p4_56p0.pt --source data/images/bus.jpg

检测结果将保存在runs/detect/exp目录下。以下是使用预训练模型检测公交车的示例：

图3：EdgeAI-YOLOv5预训练模型检测公交车效果，成功识别出车辆和行人

🧠 深入理解：模型架构与优化

EdgeAI-YOLOv5引入了多种优化技术，其中SPP（空间金字塔池化）模块的改进尤为关键。通过使用不同大小的池化核，模型能够有效提取多尺度特征：

图4：SPP模块的三种实现方案，EdgeAI版本采用了进一步优化的池化结构，提升特征提取能力

模型配置文件位于models/目录，如yolov5s6.yaml定义了网络结构和超参数。通过修改这些文件，可以调整模型深度和宽度以适应不同需求。

📊 自定义数据集训练全流程

1️⃣ 数据集准备

按照YOLO格式组织数据集，目录结构如下：

dataset/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/

创建数据集配置文件custom_data.yaml，指定训练集和验证集路径、类别数及类别名称：

train: ../dataset/images/train val: ../dataset/images/val nc: 2 # 类别数 names: ['class1', 'class2'] # 类别名称

2️⃣ 开始训练

使用以下命令启动训练，指定数据集配置、模型配置和预训练权重：

python train.py --data custom_data.yaml --cfg models/yolov5s6.yaml --weights pretrained_models/models/detection/coco/edgeai-yolov5/yolov5s6_640_ti_lite_37p4_56p0.pt --batch-size 16 --epochs 100

训练过程中，模型权重会保存在runs/train/exp/weights/目录，包括best.pt（验证集效果最佳）和last.pt（最后一轮权重）。

3️⃣ 模型评估与优化

训练完成后，使用验证集评估模型性能：

python val.py --data custom_data.yaml --weights runs/train/exp/weights/best.pt

根据评估结果，可以调整超参数（如学习率、数据增强策略）进一步优化模型。超参数配置文件位于pretrained_models/checkpoints/detection/coco/edgeai-yolov5/目录下的hyp.yaml。

💡 实用技巧与常见问题

模型导出：使用export.py将训练好的模型导出为ONNX格式，便于部署到边缘设备：
```
python export.py --weights runs/train/exp/weights/best.pt --include onnx
```
推理速度优化：对于实时应用，可使用--img-size参数减小输入图像尺寸，如--img 320以提升速度。
数据集增强：修改utils/augmentations.py文件，添加自定义数据增强方法，提高模型泛化能力。