[目标检测] 如何获取数据集对应的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))}}};")
相关推荐
昨夜见军贴06165 分钟前
IACheck × AI审核重构检测方式:破解工业检测报告频繁返工的根本难题
人工智能·重构
知乎的哥廷根数学学派22 分钟前
基于自适应多尺度小波核编码与注意力增强的脉冲神经网络机械故障诊断(Pytorch)
人工智能·pytorch·python·深度学习·神经网络·机器学习
好奇龙猫1 小时前
【AI学习-comfyUI学习-三十二节-FLXU原生态反推+controlnet depth(UNion)工作流-各个部分学习】
人工智能·学习
wyw00001 小时前
目标检测之Faster R-CNN
计算机视觉
peixiuhui1 小时前
EdgeGateway 快速开始手册-表达式 Modbus 报文格式
人工智能·mqtt·边缘计算·iot·modbus tcp·iotgateway·modbus rtu
bing.shao2 小时前
golang 做AI任务执行
开发语言·人工智能·golang
鼎道开发者联盟2 小时前
2025中国AI开源生态报告发布,鼎道智联助力产业高质量发展
人工智能·开源·gui
贾维思基2 小时前
告别RPA和脚本!视觉推理Agent,下一代自动化的暴力解法
人工智能·agent
P-ShineBeam2 小时前
引导式问答-对话式商品搜索-TRACER
人工智能·语言模型·自然语言处理·知识图谱