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 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