How does onnx change the dimension of the input

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.
insert image description here
Let's load the model first and look at its input and node.
insert image description here
insert image description here
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
insert image description here

Guess you like

Origin blog.csdn.net/qq_16792139/article/details/131812333