地平线旭日x3派·yolov8

本文介绍了如何将yolov8的预训练模型从PyTorch转换为ONNX格式,以便在地平线BPU上进行推理。步骤包括下载预训练权重、修改head.py中的Detect类以及使用export.py脚本进行模型转换。
摘要由CSDN通过智能技术生成

官方预训练模型转换

  1. 下载yolov8源码解压到本地,并配置基础运行环境。
  2. 下载官方预训练模型
  1. 进入ultralytics-main目录下,新建文件夹weights,并将步骤2中下载的权重文件放进去。
  2. 修改ultralytics/nn/modules/head.py文件内容->Detect类
   def forward(self, x):
        results = []
        for i in range(self.nl):
            dfl = self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()
            cls = self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()
            results.append(torch.cat([cls, dfl], -1))
        return tuple(results)
  1. 新建export.py文件
#encoding:utf-8
from ultralytics import YOLO
import onnx
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./weights/yolov8n.pt', help='initial weights path') 
#================================================================
opt = parser.parse_args()
print(opt)
    
model_path = opt.weights
# load a pretrained model
model = YOLO(model_path)

# export onnx
success = model.export(format='onnx', opset=11, simplify=True, dynamic=False, imgsz=640)
assert success

model = onnx.load(model_path.replace(".pt",".onnx"))

# 修改输入输出张量的名称
idx_start = 0
for input in model.graph.input:
    for node in model.graph.node:
        # 如果当前节点的输入名称与待修改的名称相同,则将其替换为新名称
        for i, name in enumerate(node.input):
            if name == input.name:
                node.input[i] = "images"
    input.name = "images"
    idx_start += 1

idx_start = 0
for output in model.graph.output:
    for node in model.graph.node:
        # 如果当前节点的输入名称与待修改的名称相同,则将其替换为新名称
        for i, name in enumerate(node.output):
            if name == output.name:
                node.output[i] = "output" + str(idx_start)
    output.name = "output" + str(idx_start)
    idx_start += 1

# 保存修改后的模型
onnx.save(model, model_path.replace(".pt",".onnx"))
  1. 命令行执行python3 export.py脚本(默认为yolov8n.pt, 加–weights参数可指定权重),转换成功会输出一下信息, 转换后的模型存于权重同级目录(*.onnx后缀模型) 注意opset_version必须为10/11
ONNX: starting export with onnx 1.14.1 opset 11...
ONNX: simplifying with onnxsim 0.4.33...
ONNX: export success ✅ 1.1s, saved as 'weights/yolov8n.onnx' (12.2 MB)

Export complete (2.6s)
Results saved to /home/zz/workspace/course/nv/example/yolov8/ultralytics-main/weights

请添加图片描述

地平线BPU加载推理篇

  • 使用地平线BPU SDK接口加载推理模型, 后续会放出代码。
    请添加图片描述
    请添加图片描述
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值