参考文件
https://download.csdn.net/download/qq_31764341/88479978?spm=1001.2014.3001.5501
操作前后的onnx发出来给大家参考下
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 onnx
onnx_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]=n
return member_map
x = {
"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")
运行python 脚本 删除层
在脚本同路径拿到结果 del_rknn.onnx
结果转换
剩下的千篇一律了,找到官方的RKNN文档随便下载一下看看,最快的方案是使用docker那个拉取一下,
然后Python test.py
需要注意的是 广为流传的demo里面的test.py导出的是rk3568还是rk3566的,需要config里面配置一下targetplatform