rknn 模型要求

硬件限制让腹，在onnx模型中需要移除Detect層的操作畦贸，將最后的卷積層的輸出作為onnx模型的輸出

直接修改已有的onnx模型

可以通過直接操作onnx模型的方法手動刪除onnx模型中最后的三個(gè)卷積層后面的所有操作谤碳，并將卷積層的輸出更改為onnx模型的輸出

更改前截圖

onnx.png

刪除思路

1. 根據(jù)名字找到三個(gè)卷積層的index
2. 記錄這三個(gè)卷積層后面所有節(jié)點(diǎn)的index， 保存在數(shù)組中
3. 從大到小遍歷數(shù)組牍戚，刪除節(jié)點(diǎn)
4. 刪除多余輸出
5. 更改輸出 （名字和輸出數(shù)據(jù)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

模型轉(zhuǎn)換主要代碼
全部代碼可以在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')