onnx如何改变输入的维度

最近遇到一个难题,就算在用行为识别onnx转rknn的时候提示维度不对,因为行为识别模型是5维的。而rknn只支持4维。
在这里插入图片描述
我们先加载模型看一下它的input和node
在这里插入图片描述
在这里插入图片描述
可以看出模型的input[1]是一个全连接,因此我们可以直接修改他的input[0]

input = helper.make_tensor_value_info('input', TensorProto.FLOAT, [24, 3, 256, 256])

另外我们要修改模型node[1]的输入,因为之前node[1]的input是原先的input[0]它的name为onnx::Reshape_0。因此我们新建一个Reshape节点。

reshape_node = helper.make_node(

    'Reshape',  # 节点类型

    inputs=['input',"/Constant_output_0"],  # 输入张量列表

    outputs=['/Reshape_output_0']  # 输出张量列表

)

同时不能忘记之前的Reshape节点的input输入有两个,还有一个"/Constant_output_0"
}
然后我们替换之前的input和node,并检验模型的有效性,最后输出。

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.!")

改完后的结构为
在这里插入图片描述

猜你喜欢

转载自blog.csdn.net/qq_16792139/article/details/131812333