从零开始训练YOLOv5:数据标注到模型导出
在智能制造车间的一条SMT贴片生产线上,每秒都有数十块PCB板经过视觉检测工位。传统基于规则的图像处理方法面对日益复杂的元器件布局和新型缺陷类型时显得力不从心——划痕、偏移、漏贴等异常难以统一建模,每次产品换线都要重新调试算法参数。这时,一个能“看懂”图像内容的AI模型就成了破局关键。
YOLOv5正是这样一套被广泛验证的解决方案。它不只是一个目标检测算法,更是一整套从数据准备、模型训练到部署落地的完整工具链。相比学术研究中常见的复杂框架,YOLOv5真正做到了让工程师用几行命令就能把深度学习模型投入实际运行。
这套系统的核心在于其高度工程化的架构设计。以Ultralytics开源实现为基础,YOLOv5采用PyTorch构建,支持端到端的训练与推理流程。它的主干网络CSPDarknet通过跨阶段部分连接(Cross-Stage Partial Connections)有效缓解了深层网络中的梯度消失问题,而PANet(Path Aggregation Network)结构则增强了高低层特征之间的信息流动,使得小目标也能被稳定捕捉。
有意思的是,尽管名为“v5”,但它并非由原始YOLO作者Joseph Redmon发布,而是Ultralytics团队独立开发的一套优化实现。这反倒让它摆脱了纯学术导向的束缚,在实用性上走得更远。例如,AutoAnchor机制会在训练前自动分析你的数据集中目标的宽高分布,生成最匹配的锚框尺寸——这意味着你不需要手动调整先验框,尤其适合工业质检这类特定领域场景。
整个检测流程可以概括为四个步骤:首先将输入图像缩放到固定尺寸(如640×640),然后由CSPDarknet提取多尺度特征;接着通过PANet进行特征融合;最后在三个不同分辨率的特征图上并行预测边界框和类别概率。由于是单阶段检测器,整个过程无需区域建议网络(RPN),一次前向传播即可完成所有对象的定位与分类,推理速度远超Faster R-CNN等两阶段方案。
# train.py - 训练入口脚本 import torch from models.common import DetectMultiBackend from utils.dataloaders import create_dataloader from utils.train import train # 配置参数 data = 'data/coco.yaml' # 数据集配置文件 cfg = 'models/yolov5s.yaml' # 模型结构定义 weights = 'yolov5s.pt' # 预训练权重路径 epochs = 100 batch_size = 16 imgsz = 640 # 输入图像大小 # 创建数据加载器 train_loader, dataset = create_dataloader( path='dataset/images/train', imgsz=imgsz, batch_size=batch_size, augment=True, # 启用Mosaic等增强 hyp='data/hyps/hyp.scratch-low.yaml' ) # 加载模型 model = DetectMultiBackend(weights, device=torch.device('cuda'), fuse=True) # 开始训练 train( model=model, train_loader=train_loader, epochs=epochs, data=data, hyp='data/hyps/hyp.scratch-low.yaml', optimizer='SGD', amp=True # 混合精度训练 )这段代码展示了标准训练流程的关键环节。其中DetectMultiBackend是一个巧妙的设计,它能统一加载PyTorch、ONNX甚至TensorRT格式的权重,极大提升了模块复用性。fuse=True启用卷积与BN层融合,在推理时可减少约20%的计算量;而amp=True开启自动混合精度训练,不仅节省显存,还能加快收敛速度——这对资源有限的开发者来说非常友好。
再来看推理阶段:
# detect.py - 推理脚本片段 from models.common import DetectMultiBackend import torch import cv2 # 加载模型 model = DetectMultiBackend('yolov5s.pt', device=torch.device('cuda'), dnn=False) model.eval() # 图像预处理 img = cv2.imread('test.jpg') img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) / 255.0 img = torch.nn.functional.interpolate(img, size=640) # 前向推理 pred = model(img)[0] # NMS后处理 from utils.general import non_max_suppression pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45) # 输出结果 for det in pred: if len(det): print(f"检测到 {len(det)} 个目标:") for *xyxy, conf, cls in det: print(f"类别: {int(cls)}, 置信度: {conf:.3f}, 位置: {xyxy}")这里有几个容易忽略但至关重要的细节。一是输入图像必须归一化到[0,1]区间,并保持RGB顺序;二是输出的坐标格式为[x1,y1,x2,y2],可以直接用于OpenCV绘图;三是NMS阈值的选择需要结合业务需求权衡——过高的IoU阈值可能导致相邻目标被误删,尤其是在密集排列的PCB元件检测中。
当模型训练完成后,下一步就是导出以便部署:
# export.py - 模型导出脚本 from export import run run( weights='runs/train/exp/weights/best.pt', # 最佳权重路径 imgsz=(640, 640), batch_size=1, include=['onnx'], # 导出ONNX格式 device='cuda' )这个简单的调用背后其实完成了复杂的图优化工作。导出的ONNX模型可以在Windows/Linux服务器上通过ONNX Runtime高效运行,也可以进一步转换为TensorRT引擎部署到Jetson系列边缘设备。值得一提的是,YOLOv5还支持动态轴设置,允许输入尺寸在一定范围内变化,这对于处理不同分辨率相机流的应用特别有用。
在一个典型的PCB缺陷检测项目中,完整的实施路径通常是这样的:
- 使用LabelImg对电阻、电容、缺件三类样本进行标注,生成YOLO格式的
.txt标签文件; - 编写数据配置
my_pcb.yaml,明确训练集路径和类别名称; - 调用
train.py启动训练,过程中观察loss曲线是否平稳下降; - 用
val.py评估验证集mAP,确认无明显过拟合; - 导出ONNX模型并集成至产线软件,通过OpenCV读取相机流实时推理。
train: dataset/images/train val: dataset/images/val nc: 3 # 类别数:电阻、电容、缺件 names: ['resistor', 'capacitor', 'missing']这种流程看似简单,但在实际落地时仍有不少坑需要注意。比如输入分辨率的选择:虽然默认640×640适用于大多数场景,但如果检测目标极小(如0402封装元件),可能需要提升到1280×1280,但这会显著增加对算力的要求。再比如模型尺寸的权衡——yolov5s仅有7.2M参数,适合嵌入式设备;而追求更高精度时可选用yolov5x(87.7M参数),但需配备高性能GPU。
另一个常被低估的因素是数据质量。我们曾遇到一个案例:客户采集的训练图像全部来自同一角度,导致模型在产线实际运行时无法识别倾斜放置的PCB板。因此建议在采集阶段就覆盖各种光照条件、旋转角度和遮挡情况,每类至少准备500张高质量标注图像。此外,启用Mosaic和MixUp数据增强能有效提升泛化能力,特别是在样本不足的情况下。
| 考量项 | 推荐做法 |
|---|---|
| 输入分辨率选择 | 优先使用640×640;若目标极小且密集,可尝试1280×1280(需更强算力) |
| 模型尺寸权衡 | 边缘设备选s或m;服务器端可用l/x追求更高精度 |
| 数据质量保障 | 标注需覆盖各种光照、角度、遮挡情况;建议每类至少500张高质量标注图像 |
| 过拟合防范 | 启用Dropout、MixUp增强;监控验证集loss是否持续下降 |
| 部署格式选择 | Windows/Linux服务端 → ONNX + ONNX Runtime;嵌入式设备 → TensorRT 或 CoreML |
实践中还有一个实用技巧:开启--cache选项缓存预处理后的数据,可避免每次训练都重复执行图像变换,大幅提升迭代效率。对于内存受限的设备,还可使用--half启用半精度推理,进一步降低资源占用。
回到最初的问题——为什么越来越多的工厂选择YOLOv5?答案或许就在于它把复杂的深度学习技术包装成了“开箱即用”的工具。你不需要成为CV专家,也能搭建起可靠的视觉检测系统。无论是初创团队做原型验证,还是大型企业推进数字化转型,这套方案都展现出了惊人的适应性和落地效率。
随着YOLO系列持续演进(如YOLOv8已原生支持实例分割与姿态估计),我们可以预见,未来的工业视觉系统将不再局限于简单的“有没有”,而是能够理解“是什么状态”“朝向如何”“是否变形”等更深层次的信息。而这一切的起点,往往就是一次成功的从零训练。
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考