官方预训练模型转换
- 下载yolov8源码解压到本地,并配置基础运行环境。
- 下载官方预训练模型
- 进入ultralytics-main目录下,新建文件夹weights,并将步骤2中下载的权重文件放进去。
- 修改ultralytics/nn/modules/head.py文件
def forward(self, x):
"""Perform forward pass through YOLO model and return predictions."""
if torch.onnx.is_in_onnx_export():
results = []
for i in range(self.nl):
results.append(nn.Flatten()(self.cv2[i](x[i]).permute(0, 2, 3, 1).contiguous()))
results.append(nn.Flatten()(self.cv3[i](x[i]).permute(0, 2, 3, 1).contiguous()))
return results
- 新建export_nnie.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=12, 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] = "data"
input.name = "data"
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] = "out" + str(idx_start)
output.name = "out" + str(idx_start)
idx_start += 1
# 保存修改后的模型
onnx.save(model, model_path.replace(".pt",".onnx"))
- 命令行执行python3 export_nnie.py脚本(默认为yolov8n.pt, 加–weights参数可指定权重),转换成功会输出一下信息, 转换后的模型存于权重同级目录(*-op.onnx后缀模型)
Namespace(weights='./weights/yolov8n.pt')
finished exporting onnx
海思相机植入-模型转换篇
前期准备
- Hi3516/3519芯片系列软件定义相机
- Aistudio开发工具 用于转换NNIE模型
- vscode/MobaXterm 用于相机后台调试
- pycaffe环境 支持上采样、permute等自定义操作
详细流程
-
onnx->caffe
可参考yolov5_onnx2caffe开源工程, 转换成功会输出一下信息。
-
打开Aistudio开发工具,新建nnie项目, 主要soc version要与目标相机版本相匹配
-
打开.cfg文件,配置onnx2nnie必要参数, 然后点击运行按钮
-
运行成功会输出以下信息
海思相机植入-NNIE加载推理篇
- 使用海思NNIE服务接口加载推理模型(yolov8n)
下集预告
- 海思NNIE开发:yolov5-face