深度学习(修改onnx文件batchsize)
如果从网上down下来的onnx文件的batchsize不是1或者是动态batchsize。
下面代码可以强行将batchsize设为1,方便推理。
代码如下:
import onnx # 1. 加载原始模型 model = onnx.load("feature.onnx") # 2. 修改输入层batch size for input in model.graph.input: dim = input.type.tensor_type.shape.dim dim[0].dim_value = 1 # 3. 修改输出层batch size for output in model.graph.output: dim = output.type.tensor_type.shape.dim dim[0].dim_value = 1 # 4. 修改所有ValueInfo的batch size for value_info in model.graph.value_info: dim = value_info.type.tensor_type.shape.dim dim[0].dim_value = 1 # 5. 保存修改后的模型 onnx.save(model, "feature_bs1.onnx")