yolov5 onnx转rknn

rknn 模型要求

硬件限制，在onnx模型中需要移除Detect层的操作，将最后的卷积层的输出作为onnx模型的输出

直接修改已有的onnx模型

可以通过直接操作onnx模型的方法手动删除onnx模型中最后的三个卷积层后面的所有操作，并将卷积层的输出更改为onnx模型的输出

更改前截图

onnx.png

删除思路

1. 根据名字找到三个卷积层的index
2. 记录这三个卷积层后面所有节点的index， 保存在数组中
3. 从大到小遍历数组，删除节点
4. 删除多余输出
5. 更改输出 （名字和输出数据shape）

代码

import onnx
onnx_file = "yolov5-ming.onnx"
save = "delete_detect.onnx"
model = onnx.load(onnx_file)

node = model.graph.node

index = []
Conv_248_index = -1
Conv_298_index = -1
Conv_198_index = -1
start_index = 100009
for i in range(len(node)):
    if node[i].name == "Conv_248":
        Conv_248_index = i
    if node[i].name == "Conv_298":
        Conv_298_index = i
    if  node[i].name == "Conv_198":
        Conv_198_index = i
        start_index = i
    if i > start_index:
        if i != Conv_248_index and i != Conv_298_index:
            index.append(i)

print(index)
for i in reversed(index):
    node.remove(node[i])
out = model.graph.output
del out[0]
out[2].name = "442"
out[2].type.tensor_type.shape.dim[0].dim_value = 1
out[2].type.tensor_type.shape.dim[1].dim_value = 42
out[2].type.tensor_type.shape.dim[2].dim_value = 28
out[2].type.tensor_type.shape.dim[3].dim_value = 28

del out[2].type.tensor_type.shape.dim[4]

out[1].name = "384"
out[1].type.tensor_type.shape.dim[0].dim_value = 1
out[1].type.tensor_type.shape.dim[1].dim_value = 42
out[1].type.tensor_type.shape.dim[2].dim_value = 56
out[1].type.tensor_type.shape.dim[3].dim_value = 56

del out[1].type.tensor_type.shape.dim[4]


out[0].name = "326"
out[0].type.tensor_type.shape.dim[0].dim_value = 1
out[0].type.tensor_type.shape.dim[1].dim_value = 42
out[0].type.tensor_type.shape.dim[2].dim_value = 112
out[0].type.tensor_type.shape.dim[3].dim_value = 112

del out[0].type.tensor_type.shape.dim[4]
onnx.save(model, save)
onnx.checker.check_model(model)

更改后截图

out1.png

out2.png

out3.png

模型转换主要代码
全部代码可以在rknn的github中找到

# Create RKNN object
    rknn = RKNN(verbose=True)

    # pre-process config
    print('--> Config model')
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]])
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export rknn model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')