Recently, I encountered a problem, even if the prompt dimension is wrong when using behavior recognition onnx to rknn, because the behavior recognition model is 5-dimensional. And rknn only supports 4 dimensions.
Let's load the model first and look at its input and node.
We can see that the input[1] of the model is a full connection, so we can directly modify its input[0]
input = helper.make_tensor_value_info('input', TensorProto.FLOAT, [24, 3, 256, 256])
In addition, we need to modify the input of the model node[1], because the previous input of node[1] is the original input[0] whose name is onnx::Reshape_0
. So we create a new Reshape node.
reshape_node = helper.make_node(
'Reshape', # 节点类型
inputs=['input',"/Constant_output_0"], # 输入张量列表
outputs=['/Reshape_output_0'] # 输出张量列表
)
At the same time, we must not forget that the previous Reshape node has two inputs and one "/Constant_output_0"
}
, then we replace the previous input and node, and check the validity of the model, and finally output.
model.graph.input[0].CopyFrom(input)
model.graph.node[1].CopyFrom(reshape_node)
# 修改权重后储存成新的onnx,不然不生效
onnx.checker.check_model(model)
onnx.save_model(model, "new_tsn.onnx")
print("Done.!")
The modified structure is