深度学习(修改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")
posted @ 2025-08-17 12:39  Dsp Tian  阅读(38)  评论(0)    收藏  举报