当前位置: 首页 > 工具软件 > yolov5 > 使用案例 >

yolov5中anchors的保存/读取

田谦
2023-12-01

yolov5中anchors的保存/读取

训练样本不同时,yolov5可能会重新计算anchors,这些anchors将保存在模型权重文件中,直接使用detect.py输出时会自动加载计算,但在某些任务(如进行板端后处理过程中),还是需要获取到这些anchors值。

如何快速查看

可以运行detect.py,运行到pre=model()前,查看model.model.state_dict()的值,可以找到anchors信息;早期版本yolov5中也存储了anchor_grid,v6版本目前只有anchor,anchor_grid是相对原图的anchor大小,anchor是相对最后的特征图的大小,所以二者的值之间相差了8,16和32倍(下采样倍数)。

从权重文件中快速提取

import torch
import sys
sys.path.append("path/yolov5-master")
weights = 'yolov5s.pt'
model = torch.load(str(weights[0] if isinstance(weights, list) else weights), map_location='cpu')
model1 = model['ema' if model.get('ema') else 'model']
model2 = model1.float().fuse().model.state_dict()

for k,v in model2.items():
    if 'anchor' in k:
        # print(k)
        # print(v)
        print(v.numpy().flatten().tolist())

打印的值即为anchor_grid/anchor的值

 类似资料: