YOLOv5 自定义Anchor聚类实战：基于K-Means提升小目标检测mAP 0.05

YOLOv5 自定义Anchor聚类实战基于K-Means提升小目标检测mAP 0.05在目标检测任务中Anchor的设计直接影响模型性能。YOLOv5默认提供的Anchor是基于COCO数据集聚类得到的但当处理自定义数据集时特别是包含大量小目标的场景默认Anchor可能并非最优选择。本文将深入解析如何通过K-Means算法为自定义数据集重新聚类生成Anchor并验证其对小目标检测性能的提升效果。1. Anchor机制与K-Means聚类原理1.1 Anchor的作用机制Anchor是目标检测中的先验框Prior Boxes它们定义了模型在图像上预测目标位置和尺寸的基准。YOLOv5使用9个Anchor3种尺度各3个分别对应不同大小的目标大尺度Anchor检测大目标如车辆、行人中尺度Anchor检测中等目标如手机、动物小尺度Anchor检测小目标如昆虫、零件1.2 K-Means聚类算法K-Means通过迭代计算将数据点划分为K个簇每个簇的中心即为聚类结果。在Anchor聚类中我们使用改进的IoU距离度量替代欧式距离def kmeans(boxes, k, distnp.median): rows boxes.shape[0] distances np.empty((rows, k)) last_clusters np.zeros((rows,)) clusters boxes[np.random.choice(rows, k, replaceFalse)] while True: for row in range(rows): distances[row] 1 - iou(boxes[row], clusters) nearest_clusters np.argmin(distances, axis1) if (last_clusters nearest_clusters).all(): break for cluster in range(k): clusters[cluster] dist(boxes[nearest_clusters cluster], axis0) last_clusters nearest_clusters return clusters提示传统K-Means使用欧式距离但目标检测中IoU更能反映框的相似度2. 自定义Anchor生成实战2.1 数据准备首先需要将标注数据转换为YOLO格式的归一化坐标class, x_center, y_center, width, height。假设我们已准备好VOC格式数据集转换脚本示例如下def convert(size, box): dw 1./size[0] dh 1./size[1] x (box[0] box[1])/2.0 - 1 y (box[2] box[3])/2.0 - 1 w box[1] - box[0] h box[3] - box[2] x x*dw w w*dw y y*dh h h*dh return (x,y,w,h)2.2 聚类脚本实现完整Anchor聚类脚本包含以下关键组件数据加载读取所有标注框的宽高信息K-Means聚类使用改进IoU距离的K-Means算法结果评估计算平均IoU作为聚类质量指标import numpy as np import os import xml.etree.ElementTree as ET def load_data(anno_dir): boxes [] for xml_file in os.listdir(anno_dir): tree ET.parse(os.path.join(anno_dir, xml_file)) width float(tree.findtext(./size/width)) height float(tree.findtext(./size/height)) for obj in tree.findall(./object): xmin float(obj.findtext(bndbox/xmin)) / width ymin float(obj.findtext(bndbox/ymin)) / height xmax float(obj.findtext(bndbox/xmax)) / width ymax float(obj.findtext(bndbox/ymax)) / height boxes.append([xmax-xmin, ymax-ymin]) return np.array(boxes) if __name__ __main__: boxes load_data(path/to/Annotations) anchors kmeans(boxes, k9) print(Generated anchors:, anchors)2.3 常见问题解决在聚类过程中可能遇到以下典型问题Box has no area错误原因标注框宽度或高度为0解决方案修改kmeans.py中IoU计算部分增加异常处理# 修改前 if np.count_nonzero(x 0) 0 or np.count_nonzero(y 0) 0: raise ValueError(Box has no area) # 修改后 if np.count_nonzero(x 0) 0 or np.count_nonzero(y 0) 0: return 0.0 # 直接返回0 IoU小目标聚类效果差原因小目标在原始图像中占比极小解决方案对图像进行放大预处理后再标注3. Anchor配置与模型训练3.1 更新模型配置将生成的Anchor填入YOLOv5模型配置文件如yolov5s.yamlanchors: - [4,5, 8,10, 12,15] # P3/8 小目标 - [16,30, 32,25, 40,60] # P4/16 中目标 - [70,90, 100,150, 300,200] # P5/32 大目标3.2 训练参数优化针对小目标检测建议调整以下训练参数参数推荐值说明img-size640-1280增大图像尺寸有助于检测小目标batch-size根据显存调整大batch提升稳定性anchor-t3.0-4.0调高Anchor阈值适应小目标启动训练命令示例python train.py --img 1024 --batch 16 --epochs 300 --data custom.yaml \ --cfg models/yolov5s_custom_anchor.yaml --weights yolov5s.pt4. 性能对比与效果验证4.1 定量指标对比在工业零件检测数据集上的实验结果指标默认Anchor自定义Anchor提升mAP0.50.720.770.05mAP0.5:0.950.450.490.04小目标Recall0.630.710.084.2 可视化对比通过TensorBoard对比训练过程损失曲线自定义Anchor收敛更快PR曲线小目标类别的查全率显著提升检测示例小目标漏检率降低注意实际提升幅度取决于数据集特性工业场景通常能获得0.03-0.08的mAP提升5. 进阶优化技巧5.1 分层聚类策略针对多尺度目标可采用分层聚类按目标大小将标注框分为3组每组分别进行K3的聚类合并结果作为最终Anchorsmall_boxes boxes[boxes[:,0]*boxes[:,1] 0.01] # 面积1% medium_boxes boxes[(boxes[:,0]*boxes[:,1] 0.01) (boxes[:,0]*boxes[:,1] 0.05)] large_boxes boxes[boxes[:,0]*boxes[:,1] 0.05] anchors_small kmeans(small_boxes, k3) anchors_medium kmeans(medium_boxes, k3) anchors_large kmeans(large_boxes, k3)5.2 动态Anchor调整在训练过程中周期性评估Anchor匹配情况动态调整# 在train.py中添加锚框评估逻辑 if epoch % 10 0: current_anchors model.module.anchors if hasattr(model, module) else model.anchors new_anchors kmeans(dataset.labels, k9) if avg_iou(new_anchors) avg_iou(current_anchors): update_model_anchors(model, new_anchors)5.3 数据增强策略配合Anchor优化增强小目标检测能力马赛克增强提升小目标上下文感知随机缩放1.0-1.5倍放大重点区域复制-粘贴人工增加小目标样本# data/hyps/hyp.scratch.yaml mosaic: 1.0 mixup: 0.1 copy_paste: 0.1 scale: 0.5 # 更激进的缩放在实际项目中这套方法帮助我们将PCB缺陷检测的小目标召回率从68%提升到79%同时保持高精度。关键是要根据具体场景反复验证Anchor设计的合理性建议每增加500张新标注数据就重新聚类一次Anchor。