[目标检测] 如何获取数据集对应的anchor size

背景:由于yolov5需要根据size生成候选anchor,因此用yolov5训练某个特定的数据集时,如果没有设置对应的anchor size。那么可能出现很多乱框的现象。

可以通过.pt文件就能读取到anchor size的参数。输入是.pt文件,输出是

python 复制代码
const int anchor0[6] = {3, 5, 6, 5, 5, 11};
const int anchor1[6] = {10, 9, 10, 17, 21, 11};
const int anchor2[6] = {20, 25, 39, 21, 47, 47};
python 复制代码
import torch
from models.experimental import attempt_load

model = attempt_load('/root/yolov5-master/runs/train/exp8/weights/best.pt')
m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1]

# Convert to CPU for easier manipulation
for i, grid in enumerate(m.anchor_grid):
    grid = grid.cpu()  # Move to CPU if on GPU
    anchor_sets = []

    # Traverse the second dimension (3 anchor sets)
    for j in range(3):  # Since the second dimension is of size 3
        # Extract the first two 'anchors' from each set (assuming we only want the first two for each set)
        anchor_pair = grid[0, j, 0, 0].tolist()  # [0, 0] for spatial dimensions, assuming we want the first anchor
        anchor_sets.append([round(x) for x in anchor_pair])

    # Flatten and format for C-style declaration
    flat_anchors = [item for pair in anchor_sets for item in pair]
    print(f"const int anchor{i}[6] = {{{', '.join(map(str, flat_anchors))}}};")
相关推荐
MonkeyKing_sunyuhua41 分钟前
6.5 行业特定应用:金融、医疗、制造等行业的定制化解决方案
人工智能·agent
god_Zeo1 小时前
从头训练小模型: 4 lora 微调
人工智能·机器学习
开心的AI频道2 小时前
GPT-4o 图像生成与八个示例指南
人工智能
%d%d22 小时前
RuntimeError: CUDA error: __global__ function call is not configured
人工智能·深度学习·机器学习
白熊1882 小时前
【计算机视觉】pyrealsense2:Intel RealSense 深度相机的 Python 接口实战指南
python·数码相机·计算机视觉
阿维的博客日记2 小时前
ϵ-prediction和z0-prediction是什么意思
人工智能·深度学习·机器学习
学术交流3 小时前
2025年软件工程与数据挖掘国际会议(SEDM 2025)
论文阅读·人工智能·数据挖掘·软件工程·论文笔记
生信漫谈3 小时前
Rice Science∣武汉大学水稻研究团队发现水稻壁相关激酶OsWAKg16和OsWAKg52同时调控水稻抗病性和产量
人工智能·学习方法
TO ENFJ4 小时前
day 10 机器学习建模与评估
人工智能·机器学习
高效匠人4 小时前
文章五《卷积神经网络(CNN)与图像处理》
图像处理·人工智能·python·cnn