rknn 模型要求
硬件限制，在onnx模型中需要移除Detect层的操作，将最后的卷积层的输出作为onnx模型的输出
直接修改已有的onnx模型
可以通过直接操作onnx模型的方法手动删除onnx模型中最后的三个卷积层后面的所有操作，并将卷积层的输出更改为onnx模型的输出
更改前截图
onnx.png
删除思路
- 根据名字找到三个卷积层的index
- 记录这三个卷积层后面所有节点的index， 保存在数组中
- 从大到小遍历数组，删除节点
- 删除多余输出
- 更改输出 （名字和输出数据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')
