最近遇到一个难题,就算在用行为识别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.!")
改完后的结构为