修改onnx batchsize

 
import onnx

def change_input_dim(model, ):
    batch_size = "16"

    # The following code changes the first dimension of every input to be batch_size
    # Modify as appropriate ... note that this requires all inputs to
    # have the same batch_size
    inputs = model.graph.input
    for input in inputs:
        # Checks omitted.This assumes that all inputs are tensors and have a shape with first dim.
        # Add checks as needed.
        dim1 = input.type.tensor_type.shape.dim[0]
        # update dim to be a symbolic value
        if isinstance(batch_size, str):
            # set dynamic batch size
            dim1.dim_param = batch_size
        elif (isinstance(batch_size, str) and batch_size.isdigit()) or isinstance(batch_size, int):
            # set given batch size
            dim1.dim_value = int(batch_size)
        else:
            # set batch size of 1
            dim1.dim_value = 1

def apply(transform, infile, outfile):
    model = onnx.load(infile)
    transform(model, )
    onnx.save(model, outfile)

apply(change_input_dim, './model/VehicleTypeClassify_0_0_0_1.onnx',
      './model/VehicleTypeClassify_0_0_0_1_16.onnx')