yolo5 onnx2rknn 瑞芯微香橙派 rk3588

本文主要是介绍yolo5 onnx2rknn 瑞芯微香橙派 rk3588，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

yolo 训练

我用的环境是yolo5.6.0 应该是
然后使用同环境下的export 注意不要换环境。。。

python export.py --weights D:\project\Pythonproj\yolov5\yolo5\runs\train\exp5\weights\best.pt --img 640 --batch 1 --include onnx --opset 12

导出 onnx ,如果你和我的版本完全相同，那么你的onnx路径应该是
D:\project\Pythonproj\yolov5\yolo5\runs\train\exp5\weights\best.onnx

yolo onnx 剪裁

原则上,如果你自己写后处理函数/对yolo后处理函数非常熟悉的话，你可以直接改他的后处理，应该也是可以的
那么如果你是新手，并且想快速的在3588上面部署yolo5 看下面
使用onnx可视化网站（网不好的话可以fanqiang 加速） https://netron.app/ 查看你的生成模型的可视化结果

看到了之后，把transpose前面的卷积层（三个卷积层的名字找到）从左到右分别是1 2 3 ，分别放到下面的三个名字的位置

使用下面的脚本

import onnx
from onnx import helper, checker
from onnx import TensorProto
import re
import argparse
# model = "D:\project\Pythonproj\yolov5\yolo5\yolov5s.onnx"
model = r"D:\project\Pythonproj\yolov5\yolo5\runs\train\exp5\weights\best.onnx"
# model = r"D:\project\Pythonproj\yolov5\yolo5\runs\train\exp15\weights\best.onnx"# model = "D:\project\caffe\dockerfile\yolov5s-simple.onnx"
import onnxonnx_model = onnx.load(model)
graph = onnx_model.graph
# print(graph)
node = graph.node# node[213].output[0] = node[212].output[0]
# node[213].output[0] = node[213].input[0]
# for idx in graph.node:
#     print(idx)
# graph.node[]
print(graph.output[0].type.tensor_type.shape)# graph.output
# graph.output[1].type.tensor_type.elem_type = 1
# graph.output[2].name = "output2"
def createGraphMemberMap(graph_member_list):member_map=dict()for n in graph_member_list:member_map[n.name]=nreturn member_mapx = {"Concat_302","Reshape_301","Reshape_263","Reshape_282","Sigmoid_267","Sigmoid_286","Sigmoid_248","Split_249","Split_268","Split_287","Mul_251","Mul_257","Mul_270","Mul_276","Mul_289","Mul_295","Mul_255","Mul_261","Mul_274","Mul_280","Mul_293","Mul_299","Add_253","Add_291","Add_272","Pow_259","Pow_278","Pow_297","Concat_262","Concat_281","Concat_300"}# 我没有实际用到 不用看我
# x = {"Concat_382","Reshape_257","Reshape_319","Reshape_381","Sigmoid_245","Sigmoid_307","Sigmoid_369","Split_246","Split_308","Split_370","Mul_248","Mul_253","Mul_310","Mul_315","Mul_372","Mul_377","Mul_251","Mul_255","Mul_313","Mul_317","Mul_375","Mul_379","Add_249","Add_311","Add_373","Pow_254","Pow_316","Pow_378","Concat_256","Concat_318","Concat_380","Transpose_322","Transpose_198","Transpose_260"}de = []
num = 0
#
node_map = createGraphMemberMap(graph.node)
output_map = createGraphMemberMap(graph.output)
graph.output.remove(output_map["output0"])
# new_output_node_names  = ["output0","output1","output2"]
output_shape_map = [[1,18,80,80],[1,18,40,40],[1,18,20,20]]
# for i in range(3):# new_nv = helper.make_tensor_value_info(new_output_node_names[i], TensorProto.FLOAT, output_shape_map[i])# graph.output.extend([new_nv])
output_map = createGraphMemberMap(graph.output)for i in range(len(graph.node)):if node[i].name in x:de.append(i)num = num+1
de.sort()
de.reverse()
# for i in range(num):
#     graph.node.remove(graph.node[de[i]])
print("graph_output:", graph.output)
for i in range(len(graph.node)):if node[i].name == "Conv_196":# 卷积层1 new_nv = helper.make_tensor_value_info(node[i].output[0], TensorProto.FLOAT, output_shape_map[0])graph.output.extend([new_nv])# node[i].output[0]="output0"if node[i].name == "Conv_215":# 卷积层2 new_nv = helper.make_tensor_value_info(node[i].output[0], TensorProto.FLOAT, output_shape_map[1])graph.output.extend([new_nv])#     node[i].output[0]="output1"if node[i].name == "Conv_234":# 卷积层3 new_nv = helper.make_tensor_value_info(node[i].output[0], TensorProto.FLOAT, output_shape_map[2])graph.output.extend([new_nv])#     node[i].output[0]="output2"onnx.checker.check_model(onnx_model)
onnx.save(onnx_model,"del_rknn.onnx")