[目标检测] 如何获取数据集对应的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))}}};")
相关推荐
yiersansiwu123d12 分钟前
AI二创的版权迷局与健康生态构建之道
人工智能
Narrastory17 分钟前
拆解指数加权平均:5 分钟看懂机器学习的 “数据平滑神器”
人工智能·机器学习
SelectDB20 分钟前
慢 SQL 诊断准确率 99.99%,天翼云基于 Apache Doris MCP 的 AI 智能运维实践
数据库·人工智能·apache
王中阳Go24 分钟前
05 Go Eino AI应用开发实战 | Docker 部署指南
人工智能·后端·go
腾讯云开发者28 分钟前
当10年架构师拿起AI:不是写不动了,是写得太快了
人工智能
小马过河R39 分钟前
RAG检索增强生成:通过重排序提升AI信息检索精准度
人工智能·语言模型
不惑_40 分钟前
通俗理解卷积神经网络
人工智能·windows·python·深度学习·机器学习
rayufo1 小时前
自定义数据在深度学习中的应用方法
人工智能·深度学习
梦帮科技1 小时前
量子计算+AI:下一代智能的终极形态?(第一部分)
人工智能·python·神经网络·深度优先·量子计算·模拟退火算法
山海青风1 小时前
藏文TTS介绍:6 MMS 项目的多语言 TTS
人工智能·python·神经网络·音视频