[Python] How to correctly change layers in onnx model and restore them in onnxruntime

I want to change onnx model and then restore the weights when using it with onnxruntime. However, the model doesn't seem to be changed.

First, I load the existing model and change the weights based on this output. Then I save the model and try to check if the outputs are the same as in the original model.

The model could be restored (if restore_layer is True) with onnxruntime as in this example based on this output, but it doesn't seem to work, they are always the same and the model doesn't seem to be changed. The expected behavior is when change_layer = True and restore_layer = False the script should raise an error, but it doesn't.

I tried to change a part of a model weights when loading it with onnxruntime like in this example based on this issue, but it doesn't seem to work.

How can I fix it?

Here is the code to reproduce.

import onnx # onnx=="1.17.0"
import torch # torch=="2.3.0"
from onnx import numpy_helper
import numpy as np # numpy==1.26.4
import torchvision.models as models # torchvision==0.18.0
from torchvision.models import ResNet18_Weights

change_layer = True
restore_layer = False
no_changed_layers = 20

model = models.resnet18(weights=ResNet18_Weights.DEFAULT)

model.eval()

input_data = torch.randn(1, 3, 224, 224)

torch.onnx.export(model, input_data, "some_model.onnx", verbose=False)


MODEL_PATH = "some_model.onnx"
_model = onnx.load(MODEL_PATH)
INTIALIZERS=_model.graph.initializer
Weight= {}
initializer_num_name = {}
for num, initializer in enumerate(INTIALIZERS):
W = numpy_helper.to_array(initializer)
Weight[num] = W
initializer_num_name[num] = initializer.name

Weight_num_name = {}
if change_layer:
for weight_zeros in range(0, no_changed_layers):
old_weight = numpy_helper.to_array(INTIALIZERS[weight_zeros])
new_weight = np.zeros_like(Weight[weight_zeros])
updated_weight = numpy_helper.from_array(new_weight)

updating_weight_name = _model.graph.initializer[weight_zeros].name
updated_weight.name = updating_weight_name
Weight_num_name[weight_zeros] = updating_weight_name
_model.graph.initializer[weight_zeros].CopyFrom(updated_weight)

onnx.save(_model, "model.onnx")


import onnxruntime
import numpy as np


if restore_layer:
options = onnxruntime.SessionOptions()
# options.log_verbosity_level=2
# options.log_severity_level=0
# options.enable_profiling=1
ortvalue_initializers = []
for num in range(0, no_changed_layers):
initializer = onnxruntime.OrtValue.ortvalue_from_numpy(
Weight[num]
)
ortvalue_initializers.append(initializer)
options.add_initializer(Weight_num_name[num], initializer)
import logging
logging.basicConfig(level=logging.DEBUG)
onnx_session3 = onnxruntime.InferenceSession("model.onnx",
sess_options=options,
providers=["CPUExecutionProvider"])
else:

onnx_session3 = onnxruntime.InferenceSession("model.onnx",
providers=["CPUExecutionProvider"])


onnx_session = onnxruntime.InferenceSession(MODEL_PATH)
input = onnx_session.get_inputs()[0]



generated_dummy_input_data = torch.randn(size=(input.shape[0],
input.shape[1],
input.shape[2],
input.shape[3]))


onnx_inputs = {input.name: generated_dummy_input_data.numpy()}
onnx_output = onnx_session.run(None, onnx_inputs)

onnx_inputs3 = {onnx_session3.get_inputs()[0].name: generated_dummy_input_data.numpy()}
onnx_output3 = onnx_session3.run(None, onnx_inputs3)

def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
x_max = np.max(x, axis=axis, keepdims=True)
tmp = np.exp(x - x_max)
s = np.sum(tmp, axis=axis, keepdims=True)
return tmp / s

onnx_model_predictions = onnx_output[0][0]
onnx_model_predictions3 = onnx_output3[0][0]
diff = np.abs(onnx_model_predictions - onnx_model_predictions3)


assert (diff > 1e-2).any(), "Big difference between logits of an original model and onnx one: \n%s\n%s"%(onnx_model_predictions[:50],
onnx_model_predictions3[:50])
print("all good")

Continue reading...