Witness output does not contain all batches only the first one

[italoaa]

My Model and Export Code

I defined a simple MLP for MNIST and exported it to ONNX:

# Define a larger MLP model for MNIST
class MNIST_MLP(nn.Module):
    def __init__(self):
        super(MNIST_MLP, self).__init__()
        self.fc1 = nn.Linear(28*28, 128) # Input: 28x28 pixels flattened
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10) # Output: 10 classes for MNIST

    def forward(self, x):
        x = x.view(-1, 28*28) # Flatten the input
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

# Create an instance of the model
model = MNIST_MLP()

# Create a dummy input (batch size 1, image size 28x28)
dummy_input = torch.randn(1, 1, 28, 28) # B, C, H, W

# Export the model to ONNX
torch.onnx.export(
    model,
    dummy_input,
    model_path,
    input_names=['input'],
    output_names=['output'],
    opset_version=11
)

print(f"Model exported to {model_path}")

EZKL setup

ezkl.gen_settings(model_path) # generate settings.json
res = await ezkl.get_srs(settings_path=settings_path, srs_path=srs_path, commitment=ezkl.PyCommitments.KZG)
ezkl.compile_circuit(model_path, compiled_model_path, settings_path)
ezkl.setup(compiled_model_path, vk_path, pk_path, srs_path)

Generate Proof

# Dummy input assumed commes from the user of the endpoint
input = torch.randn(8, 1, 28, 28) # B, C, H, W
# Reshape the input to a list of flattened images
data = dict(input_data = input.detach().numpy().reshape(input.shape[0], -1).tolist())

# Serialize data into file:
json.dump(data, open(data_path, 'w'))

# Generate a witness with the data and the model
res = ezkl.gen_witness(data_path, compiled_model_path, witness_path, srs_path=srs_path)

with open(witness_path, 'r') as f:
    witness_data = json.load(f)
# Outputs
print(f"Witness generated to {witness_path}")
witness_outputs = witness_data['pretty_elements']['rescaled_outputs']
print(f"Witness outputs: {witness_outputs}")
normal_outputs = model(input).detach().numpy().reshape(-1).tolist()
print(f"Model outputs: {normal_outputs}")

# compare print len
print(f"Witness outputs length: {len(witness_outputs)}")
print(f"Model outputs length: {len(normal_outputs)}")

torch_outputs = torch.tensor(normal_outputs).view(8,10)
torch.argmax(torch_outputs, dim=1)

The output of this is the following:

Witness generated to witness.json
Witness outputs: [['0.03125', '-0.1015625', '0', '-0.1171875', '-0.0703125', '-0.1796875', '0.171875', '0.2578125', '0.078125', '0.4296875']]
Model outputs: [0.05165461450815201, -0.10043540596961975, 0.01641690358519554, -0.09567052125930786, -0.06564567238092422, -0.17707957327365875, 0.1964305192232132, 0.25893378257751465, 0.07650800794363022, 0.4348622262477875, -0.07594309747219086, -0.13848379254341125, 0.047103896737098694, 0.04273897781968117, -0.03477935492992401, -0.10576177388429642, 0.03600144013762474, 0.05549345538020134, 0.2860642671585083, 0.3272760510444641, -0.0024701450020074844, -0.2362748086452484, 0.11670715361833572, 0.03200536221265793, 0.016857996582984924, -0.047733813524246216, 0.11958855390548706, 0.08800196647644043, 0.05886254459619522, 0.31702882051467896, 0.047089021652936935, -0.08444338291883469, 0.08129534870386124, 0.012957751750946045, 0.06360675394535065, -0.06898896396160126, 0.0897778645157814, 0.07359668612480164, 0.17615553736686707, 0.20919445157051086, 0.06776668131351471, -0.14936968684196472, 0.23163698613643646, 0.13424430787563324, -0.10164189338684082, 0.039224691689014435, 0.1602347493171692, 0.1562921553850174, 0.06017901003360748, 0.3460584282875061, 0.08242134004831314, -0.14110985398292542, 0.13711898028850555, -0.06065642833709717, 0.013093873858451843, -0.046491846442222595, 0.016045652329921722, -0.03409861773252487, 0.17993345856666565, 0.11482946574687958, -0.06223057210445404, -0.27873119711875916, 0.023011811077594757, -0.05892151594161987, 0.10935309529304504, -0.046927616000175476, 0.24022577702999115, 0.02058327943086624, 0.1436556875705719, 0.11812470853328705, -0.05804021656513214, -0.12759438157081604, 0.08285728842020035, -0.06750114262104034, 0.08212511241436005, 0.01702721416950226, 0.15213808417320251, -0.07221671938896179, 0.17901360988616943, 0.27853965759277344]
Witness outputs length: 1
Model outputs length: 80
tensor([9, 9, 9, 9, 9, 8, 6, 9])

I feel like I am missing something, as you can see the outputs of the witness is just 10 elements while the outputs of the normal model are 80 (which makes sense for 8 images). Is this the expected behavior? or am I setting something wrong? Did I forget to do a step?

I also runned the example: https://github.com/zkonduit/ezkl/blob/main/examples/notebooks/mnist_classifier.ipynb
and it also appears to have the same issue... Is this expected then? If it is expected, in an inference service how should I approach processing a batch of images?

Disclaimer: I am new, sorry if I have missed something, I am trying to get my head around this.

[italoaa]

Sorry about that found the Batched example! For anyone looking for the fix the key is this:

py_run_args = ezkl.PyRunArgs()
py_run_args.variables = [("batch_size", 2)]
res = ezkl.gen_settings(model=model_path, output=settings_path, py_run_args=py_run_args)

My guess is that it links that variable name to the dynamic axis while exporting:

torch.onnx.export(model,                     # model being run
                  x_tensor,   # model input (or a tuple for multiple inputs) - Ensure input is a tensor
                  model_path,                # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
                                'output' : {0 : 'batch_size'}})