[目标检测] 如何获取数据集对应的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))}}};")
相关推荐
小和尚同志4 小时前
AI 自动化测试探索(二):Chrome-devtools MCP
人工智能·e2e·aigc
冬奇Lab6 小时前
Workflow 系列(02):设计范式——四层架构、三种 Context 传递模式与确认门设计
人工智能·agent·工作流引擎
冬奇Lab6 小时前
每日一个开源项目(第145篇):Trellis - 把项目记忆、规范和任务上下文持久化进代码仓库
人工智能·开源·资讯
有道AI情报局6 小时前
Harness即产品
人工智能·agent
罗西的思考8 小时前
机器人 / 强化学习】HIL-SERL:人类在环驱动的具身智能进化框架
人工智能·算法·机器学习
IT_陈寒9 小时前
SpringBoot自动配置的坑,我的API突然就404了
前端·人工智能·后端
笃行3509 小时前
从零到上线:用 EdgeOne Makers + CodeBuddy 搭一个「对账核对员」AI Agent
人工智能
用户6856326208699 小时前
Claude Code 乱猜字段名?我给它写了一个"数据库查询约束 Skill"
人工智能
你_好9 小时前
# 给你的产品嵌入一个「会操作界面的 AI 助手」
人工智能
ShallWeL9 小时前
【机器学习】(3)—— 线性回归:梯度下降
人工智能·机器学习