Changing Batch SIze

I have an onnx model. I just want to change the batch size of the model.

I have attached an image of a single node of the graph. The first one shows batch size = 1 and the second one shows batch size = 4. How to make this change using python code ??

[gramalingam]

Hi,

You can change the batch-size as below. Note that you can also make the batch-size symbolic (e.g, “N”) to indicate an unknown value … then you don’t need to keep changing it for every different batch-size.

import onnx

def change_input_dim(model):
    # Use some symbolic name not used for any other dimension
    sym_batch_dim = "N"
    # or an actal value
    actual_batch_dim = 4 

    # The following code changes the first dimension of every input to be batch-dim
    # Modify as appropriate ... note that this requires all inputs to
    # have the same batch_dim 
    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
        dim1.dim_param = sym_batch_dim
        # or update it to be an actual value:
        # dim1.dim_value = actual_batch_dim


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

apply(change_input_dim, r"input-file-name, r"output-file-name")

[ivan23kor]

The scripts above did not work for me for the MNIST model (need to change graph.output, graph.value_info and reshape nodes). After more scrutiny I found this post so I implemented a Python script that works for MNIST:

import onnx
import os
import struct

from argparse import ArgumentParser


def rebatch(infile, outfile, batch_size):
    model = onnx.load(infile)
    graph = model.graph

    # Change batch size in input, output and value_info
    for tensor in list(graph.input) + list(graph.value_info) + list(graph.output):
        tensor.type.tensor_type.shape.dim[0].dim_param = batch_size

    # Set dynamic batch size in reshapes (-1)
    for node in  graph.node:
        if node.op_type != 'Reshape':
            continue
        for init in graph.initializer:
            # node.input[1] is expected to be a reshape
            if init.name != node.input[1]:
                continue
            # Shape is stored as a list of ints
            if len(init.int64_data) > 0:
                # This overwrites bias nodes' reshape shape but should be fine
                init.int64_data[0] = -1
            # Shape is stored as bytes
            elif len(init.raw_data) > 0:
                shape = bytearray(init.raw_data)
                struct.pack_into('q', shape, 0, -1)
                init.raw_data = bytes(shape)

    onnx.save(model, outfile)

if __name__ == '__main__':
    parser = ArgumentParser('Replace batch size with \'N\'')
    parser.add_argument('infile')
    parser.add_argument('outfile')
    args = parser.parse_args()

    rebatch(args.infile, args.outfile, 'N')

• Works for googlenet-7.onnx inception-v2-7.onnx mnist-8.onnx resnet152-v1-7.onnx resnet50-v1-7.onnx squeezenet1.0-7.onnx vgg16-7.onnx from the onnx models repo.