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