YOLO-v5锚框聚类：K-means生成最优先验框教程-编程实验室

YOLO-v5锚框聚类：K-means生成最优先验框教程

1. 引言

1.1 YOLO-V5 简介

YOLO（You Only Look Once）是一种流行的物体检测模型，由华盛顿大学的 Joseph Redmon 和 Ali Farhadi 开发。自2015年首次发布以来，YOLO 因其在保持高精度的同时具备极快的推理速度而广受关注。YOLO-v5 是该系列中一个高效、轻量且易于部署的版本，尽管并非官方 YOLO 团队开发，但因其出色的工程实现和社区支持，已成为工业界广泛应用的目标检测工具之一。

在 YOLO-v5 中，先验框（anchor boxes）是目标检测的关键组成部分，直接影响模型对不同尺度和长宽比目标的匹配能力。默认的 anchor 设置适用于通用数据集（如 COCO），但在特定应用场景或自定义数据集中，使用 K-means 聚类算法生成适配于实际标注分布的最优 anchor 框，能显著提升检测性能。

1.2 锚框优化的价值与必要性

标准 anchor 框基于大规模数据集统计得出，可能无法很好地适应小众或领域特定的数据分布（如无人机航拍、医学图像、工业缺陷等）。通过在自有数据集上运行 K-means 聚类，可以生成更贴合真实边界框尺寸和比例的 anchor 集合，从而提高正样本匹配率、加快收敛速度，并最终提升 mAP（mean Average Precision）指标。

本文将详细介绍如何基于 YOLO-v5 架构，利用 K-means 算法从训练集标注中自动聚类生成最优 anchor 框，并提供完整可执行代码与实践建议。

2. 原理与方法设计

2.1 先验框的作用机制

在 YOLO 系列模型中，每个特征图上的网格单元会预测多个边界框，这些预设的参考框即为“先验框”（anchors）。模型不直接预测绝对坐标，而是预测相对于 anchor 的偏移量。因此，anchor 的形状越接近真实目标，所需的调整越小，学习越稳定。

YOLO-v5 在三个不同尺度的特征图上进行预测，每层分配3个 anchor，共9个 anchor 组成一组三元组，分别对应小、中、大目标。

2.2 为何选择 K-means 而非 K-means++？

传统的欧氏距离 K-means 聚类不适合 bounding box 尺寸聚类，因为其距离度量偏向大框。YOLO 提出使用IoU（交并比）作为相似性度量，更合理地反映框之间的重叠程度。具体而言，我们采用以1 - IoU作为距离函数，使得两个框重叠越多，距离越近。

公式如下：
$$ d(box, centroid) = 1 - IoU(box, centroid) $$

该策略确保聚类结果更加符合检测任务的实际需求。

2.3 聚类流程概览

整个过程可分为以下步骤：

提取训练集中所有 ground truth 边界框的宽高（w, h）
初始化 k 个聚类中心（k=9，对应 YOLO-v5 的 anchor 数量）
计算每个 bbox 到各 centroid 的距离（基于 1-IoU）
将每个 bbox 分配给最近的 centroid
更新 centroid 为所属 bbox 的加权中心（面积加权防止极端值影响）
重复 3~5 直至收敛（centroid 变化小于阈值或达到最大迭代次数）

3. 实践实现：K-means 锚框生成全流程

3.1 环境准备与依赖安装

本实验可在任意 Linux 或 Windows 环境下运行，推荐使用 Python 3.8+ 及以下库：

pip install numpy matplotlib opencv-python torch lxml

若使用 CSDN 提供的 YOLO-V5 镜像环境，已预装所需依赖，可直接进入项目目录开始操作。

首先进入 yolov5 项目路径：

cd /root/yolov5/

创建新文件generate_anchors.py开始编写聚类脚本。

3.2 数据读取：解析 XML 或 TXT 标注

假设你的数据集采用 Pascal VOC 格式（XML）或 YOLO 格式（TXT），我们需要提取所有标注框的原始宽高信息。

以下是读取 VOC XML 格式的示例代码：

import os import xml.etree.ElementTree as ET import glob def load_ground_truths(voc_annotations_dir): """加载VOC格式标注文件中的宽高信息""" bboxes = [] xml_files = glob.glob(os.path.join(voc_annotations_dir, "*.xml")) for xml_file in xml_files: tree = ET.parse(xml_file) root = tree.getroot() size = root.find('size') img_w = int(size.find('width').text) img_h = int(size.find('height').text) for obj in root.findall('object'): bbox = obj.find('bndbox') xmin = float(bbox.find('xmin').text) ymin = float(bbox.find('ymin').text) xmax = float(bbox.find('xmax').text) ymax = float(bbox.find('ymax').text) # 转换为 w, h w = xmax - xmin h = ymax - ymin bboxes.append([w, h]) return bboxes

注意：此处保留的是原始像素单位下的宽高，无需归一化。

3.3 K-means 聚类核心实现

import numpy as np def iou_distance(boxes, centroids): """计算每个box与centroids之间的1-IoU距离""" n = boxes.shape[0] k = centroids.shape[0] distance = np.empty((n, k)) for i in range(n): for j in range(k): w_inter = min(boxes[i][0], centroids[j][0]) h_inter = min(boxes[i][1], centroids[j][1]) inter_area = w_inter * h_inter union_area = boxes[i][0] * boxes[i][1] + \ centroids[j][0] * centroids[j][1] - inter_area iou = inter_area / union_area if union_area > 0 else 0 distance[i][j] = 1 - iou return distance def kmeans_anchors(bboxes, k=9, max_iters=100, tol=1e-5): """K-means聚类生成最优anchor""" bboxes = np.array(bboxes) np.random.seed(42) indices = np.random.choice(bboxes.shape[0], k, replace=False) centroids = bboxes[indices].copy() # 初始化聚类中心 for iter_idx in range(max_iters): distances = iou_distance(bboxes, centroids) labels = np.argmin(distances, axis=1) # 分配最近簇 new_centroids = np.zeros_like(centroids) for i in range(k): cluster_boxes = bboxes[labels == i] if len(cluster_boxes) == 0: continue # 使用面积加权平均更新中心，避免异常值主导 areas = cluster_boxes[:, 0] * cluster_boxes[:, 1] weighted_sum = np.sum(cluster_boxes * areas[:, None], axis=0) total_area = np.sum(areas) new_centroids[i] = weighted_sum / total_area if total_area > 0 else centroids[i] # 判断是否收敛 shift = np.linalg.norm(new_centroids - centroids) centroids = new_centroids if shift < tol: print(f"Converged after {iter_idx+1} iterations.") break # 按面积排序（从小到大） areas = centroids[:, 0] * centroids[:, 1] sorted_indices = np.argsort(areas) centroids = centroids[sorted_indices] return centroids

3.4 执行聚类并输出结果

if __name__ == "__main__": ANNOTATIONS_DIR = "/path/to/your/voc_annotations/" # 修改为实际路径 bboxes = load_ground_truths(ANNOTATIONS_DIR) print(f"Loaded {len(bboxes)} bounding boxes.") anchors = kmeans_anchors(bboxes, k=9) print("Generated Anchors (w, h):") print(np.round(anchors, 2)) # 输出用于YAML配置的格式 print("\nYAML Format:") reshaped = anchors.reshape(3, 3, 2).tolist() # [3 layers, 3 anchors per layer, (w,h)] for i, layer in enumerate(reshaped): print(f" - {[int(round(x[0])) for x in layer]}") # width print(f" - {[int(round(x[1])) for x in layer]}") # height

输出示例（模拟）：

Generated Anchors (w, h): [[12.5 18.3] [23.1 45.6] [41.2 33.8] [55.7 72.1] [89.4 61.3] [110.2 105.6] [156.3 132.4] [210.5 180.1] [280.0 260.7]] YAML Format: - [13, 23, 41] - [18, 46, 34] - [56, 89, 110] - [72, 61, 106] - [156, 210, 280] - [132, 180, 261]

3.5 替换 YOLO-v5 配置文件中的 anchors

编辑data/hyps/hyp.scratch-low.yaml或自定义.yaml文件中的anchors字段：

anchors: - [13,18, 23,46, 41,34] # P3/8 - [56,72, 89,61, 110,106] # P4/16 - [156,132, 210,180, 280,261] # P5/32

同时，在训练命令中指定该超参文件：

python train.py --img 640 --batch 16 --epochs 100 --data custom.yaml --cfg models/yolov5s.yaml --hyp data/hyps/hyp.scratch-low.yaml

4. 实践问题与优化建议

4.1 常见问题及解决方案

问题	原因分析	解决方案
聚类结果不稳定	初始中心随机导致局部最优	多次运行取平均，或使用遗传算法改进初始化
anchor 过大或过小	数据集中存在极端尺寸目标	清洗异常标注，或限制最大最小尺寸范围
训练初期 loss 波动大	新 anchor 与网络初始权重不匹配	微调学习率起始值，增加 warmup 步数

4.2 性能优化建议

数据清洗优先：去除模糊、错误标注的框，避免噪声干扰聚类。
考虑尺度均衡：若小目标占比极高，可对 bbox 按面积采样加权，防止被大框主导。
结合多阶段聚类：先粗聚类再细调，提升稳定性。
可视化对比：绘制原始 anchor 与新 anchor 分布图，辅助判断合理性。

5. 总结

5.1 技术价值总结

本文系统阐述了 YOLO-v5 中基于 K-means 的 anchor 框生成原理与工程实现。通过使用 IoU 距离度量的聚类方法，能够有效提取出自定义数据集中目标尺寸的统计规律，生成更具代表性的先验框集合。相比默认 anchor，定制化 anchor 显著提升了模型对特定场景目标的定位能力和整体检测精度。