[目标检测] 如何获取数据集对应的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))}}};")
相关推荐
冬奇Lab几秒前
SubAgent 原理深度解析:AI 系统如何通过委托实现专业化分工
人工智能·agent·claude
ZhengEnCi13 分钟前
01c-循环神经网络RNN详解
人工智能·深度学习
仙女修炼史23 分钟前
CNN的捷径学习Shortcut Learning in Deep Neural Networks
人工智能·学习·cnn
企业架构师老王27 分钟前
数据安全优先:企业级智能体私有化部署完整方案与最佳实践
人工智能·ai
北漂Zachary29 分钟前
AI训练硬件选型:GPU算力梯队全解析
人工智能·gpu算力
CS创新实验室32 分钟前
CS实验室行业报告:机器人领域就业分析报告
大数据·人工智能·机器人
才兄说33 分钟前
机器人二次开发机器狗巡检?绕行率低于10%
人工智能
QYR-分析36 分钟前
机器人物体检测/追踪传感器行业解析及市场机遇
人工智能·机器人
Pushkin.40 分钟前
数仓 AI Coding 应用实战:基于 AI 的数仓效能提升
人工智能·数仓ai提效
罗不俷1 小时前
【机器学习】(一)机器学习入门概念
人工智能·机器学习