[目标检测] 如何获取数据集对应的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))}}};")
相关推荐
chian-ocean7 分钟前
量化加速实战:基于 `ops-transformer` 的 INT8 Transformer 推理
人工智能·深度学习·transformer
那个村的李富贵7 分钟前
从CANN到Canvas:AI绘画加速实战与源码解析
人工智能·ai作画·cann
水月wwww17 分钟前
【深度学习】卷积神经网络
人工智能·深度学习·cnn·卷积神经网络
晚霞的不甘42 分钟前
CANN 在工业质检中的亚像素级视觉检测系统设计
人工智能·计算机视觉·架构·开源·视觉检测
island131444 分钟前
CANN HIXL 高性能单边通信库深度解析:PGAS 模型在异构显存上的地址映射与异步传输机制
人工智能·神经网络·架构
前端摸鱼匠1 小时前
YOLOv8 环境配置全攻略:Python、PyTorch 与 CUDA 的和谐共生
人工智能·pytorch·python·yolo·目标检测
结局无敌1 小时前
构建百年工程:cann/ops-nn 的可持续演进之道
人工智能·cann
MSTcheng.1 小时前
CANN算子开发新范式:基于ops-nn探索aclnn两阶段调用架构
人工智能·cann
renhongxia11 小时前
如何基于知识图谱进行故障原因、事故原因推理,需要用到哪些算法
人工智能·深度学习·算法·机器学习·自然语言处理·transformer·知识图谱
做人不要太理性1 小时前
CANN Runtime 运行时与维测组件:异构任务调度、显存池管理与全链路异常诊断机制解析
人工智能·自动化