import torch
import struct
# utils.torch_utils from yolov5 github: https://github.com/ultralytics/yolov5
from utils.torch_utils import select_device
# Initialize
device = select_device('cpu')
# Load model
model = torch.load('weights/yolov5s.pt', map_location=device)['model'].float() # load to FP32
model.to(device).eval()
f = open('yolov5s.wts', 'w')
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f',float(vv)).hex())
f.write('\n')
代码来源:
https://gitee.com/hanshangying/tensorrt-yolov5/tree/main#http://openjetson.com/
版权声明:本文为lgh0824原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。