静态batch转换动态batch

王航
2023-12-01

静态batch转换动态batch

import onnx
import onnx.checker
import onnx.utils
from onnx.tools import update_model_dims

model=onnx.load('sim_plate.onnx')

dim_proto0=model.graph.input[0].type.tensor_type.shape.dim[0]
dim_proto2=model.graph.input[0].type.tensor_type.shape.dim[2]
dim_proto3=model.graph.input[0].type.tensor_type.shape.dim[3]

dim_proto0.dim_param='batch_size'
#dim_proto2.dim_param='?'
#dim_proto3.dim_param='?'
dim_proto_0=model.graph.output[0].type.tensor_type.shape.dim[0]
dim_proto_0.dim_param='batch_size'
onnx.save(model,'dynamic_sim_plate.onnx')
 类似资料: