Troubleshooting onnxruntime inference - X num_dims does not match W num_dims

Question

Using the nnhash.py script found here: https://github.com/AsuharietYgvar/AppleNeuralHash2ONNX

# Copyright 2021 Asuhariet Ygvar
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
# http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied. See the License for the specific language governing
# permissions and limitations under the License.

import sys
import onnxruntime
import numpy as np
from PIL import Image

# Load ONNX model
session = onnxruntime.InferenceSession(sys.argv[1])

# Load output hash matrix
seed1 = open(sys.argv[2], 'rb').read()[128:]
seed1 = np.frombuffer(seed1, dtype=np.float32)
seed1 = seed1.reshape([96, 128])

# Preprocess image
image = Image.open(sys.argv[3]).convert('RGB')
image = image.resize([360, 360])
arr = np.array(image).astype(np.float32) / 255.0
arr = arr * 2.0 - 1.0
arr = arr.transpose(2, 0, 1).reshape([1, 3, 360, 360])

# Run model
inputs = {session.get_inputs()[0].name: arr}
outs = session.run(None, inputs)

# Convert model output to hex hash
hash_output = seed1.dot(outs[0].flatten())
hash_bits = ''.join(['1' if it >= 0 else '0' for it in hash_output])
hash_hex = '{:0{}x}'.format(int(hash_bits, 2), len(hash_bits) // 4)

print(hash_hex)

When I try to execute using: python nnhash.py ../model.onnx ../seed1.dat ../1.png

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running FusedConv node. Name:'' Status Message: X num_dims does not match W num_dims. X: {1,1280,1,1} W: {500}

I've attached what the Netron layer output looks like: https://i.imgur.com/EeVItQ2.jpeg (sending it as a link because the actual image is extremely long)

From what I can tell, towards the very bottom, there's a W:{500} right before the leafy part. I'm sure this is what is causing the issue, I'm just not sure what I need to do in order to process the input image so that it flows through the model fine.

Edit: I found an older model.onnx that appears to work fine, the last few layers differ a bit, which is I think the problem. I'm not sure what to do in order to get the script to work with the newer model.onnx file.

Old:

New:

Answer 1

I found out that there are additional transformations that are done by the vision library, so I wrote a python script in order to just modify the last couple of layers:

import onnx
from onnx import helper, shape_inference, numpy_helper, TensorProto
import numpy as np

def is_node_used(node_name, graph):
    for node in graph.node:
        if node_name in node.input:
            return True
    return False

# Load the new ONNX model
model_path = 'model.onnx'
new_model = onnx.load(model_path)

# Load the old ONNX model to use as a reference
old_model_path = 'old-model.onnx'
old_model = onnx.load(old_model_path)

# Extract the graph from the new model
graph = new_model.graph

# Remove the last layer(s) from the new model
for node in graph.node:
    if node.output[0] == 'leaf/logits':
        graph.node.remove(node)

# Create a new last layer that matches the old model's configuration
# Define the weight and bias for the new Conv layer
W_conv_last = numpy_helper.from_array(
    np.random.rand(128, 500, 1, 1).astype(np.float32), name='W_conv_last'
)
B_conv_last = numpy_helper.from_array(
    np.random.rand(128).astype(np.float32), name='B_conv_last'
)

# Add weight and bias initializers to the graph
graph.initializer.append(W_conv_last)
graph.initializer.append(B_conv_last)

# Create the new Conv node for the last layer
new_conv_last = helper.make_node(
    'Conv',
    inputs=['relu_output', 'W_conv_last', 'B_conv_last'],
    outputs=['leaf/logits'],
    kernel_shape=[1, 1],
    pads=[0, 0, 0, 0],
    strides=[1, 1],
    dilations=[1, 1],
)

# Define weights and biases for the modified convolutional layer
W_conv_mod = numpy_helper.from_array(
    np.random.rand(500, 1280, 1, 1).astype(np.float32), name='W_conv_mod'
)
B_conv_mod = numpy_helper.from_array(
    np.random.rand(500).astype(np.float32), name='B_conv_mod'
)

# Add weight and bias initializers to the graph
graph.initializer.append(W_conv_mod)
graph.initializer.append(B_conv_mod)

# Create the new Conv node for the modified layer
new_conv_mod = helper.make_node(
    'Conv',
    inputs=['stem/MobilenetV3/GAP_13/mul_1', 'W_conv_mod', 'B_conv_mod'],
    outputs=['conv_output'],
    kernel_shape=[1, 1],
    pads=[0, 0, 0, 0],
    strides=[1, 1],
    dilations=[1, 1],
)

# Create the ReLU node
new_relu = helper.make_node(
    'Relu',
    inputs=['conv_output'],
    outputs=['relu_output'],
)

# Find and remove the old layer with the input 'stem/MobilenetV3/GAP_13/mul_1'
for node in graph.node:
    if 'stem/MobilenetV3/GAP_13/mul_1' in node.input:
        graph.node.remove(node)
        break

# Insert the new Conv node for the modified layer, followed by ReLU and the last Conv node
graph.node.append(new_conv_mod)
graph.node.append(new_relu)
graph.node.append(new_conv_last)

# Remove any floating ReLU nodes
nodes_to_remove = []
for node in graph.node:
    if node.op_type == 'Relu' and not is_node_used(node.output[0], graph):
        nodes_to_remove.append(node)

for node in nodes_to_remove:
    graph.node.remove(node)

# Save the modified ONNX model
onnx.save(new_model, 'modified_new_model.onnx')

This takes the new model and makes the last few layers look like how the old model looks and saves a new model. Now it works as expected.