Onnx: How to update the weights back into the ONNX model

[buddhapuneeth]

Ask a Question

Question

Using the below code, I can read the weights (in this case the first weight) from the model into a numpy array.
In case if I modify the weight (i.e. numpy array), how to assign it back to model.

Basically, I am trying to modify the weights.

code:

import onnx
from onnx import numpy_helper
model   = onnx.load('test.onnx')
init_0 = model.graph.initializer[0]
init_np = numpy_helper.to_array(init_0)
print(init_np)
init_np[0] = 1.1
// how to modify the actual weight here?
...
// save the new model

Further information

• Relevant Area (e.g. model usage, backend, best practices, converters, shape_inference, version_converter, training, test, operators):
  I explored on from_array but not sure of complete usage

• Is this issue related to a specific model? No

[jcwchen]

Hi @buddhapuneeth,
You can try this:

import onnx
from onnx import numpy_helper
model = onnx.load('test.onnx')
# numpy value which you want to update 
numpy_array = numpy.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=float)
# convert numpy to TensorProto
tensor = numpy_helper.from_array(numpy_array)
# copy this TensorProto to target tensor
model.graph.initializer[0].CopyFrom(tensor)
print(model.graph.initializer[0])
// then save the modified model

[buddhapuneeth]

Hi @jcwchen
Thanks for your reply. I tried based on your code. It is working, but raw_data format is changed like below:

Left: After updating, Right: Before updating
Visualizers like Neutron were not able to show the content like:

Left: Before updating, Right: After updating
I am not sure, will it cause any issue while reading using ONNX C++ APIs.
Also, in one of the other model, for LSTM in Netron I am seeing some different info for variables like 'sequence_len'. Due to large size, I am not able to convert to .txt and examine the change.

Left: Before updating, Right: After updating

def weight_update(onnx_model_path, onnx_updated_model_path, gradient_map, learning_rate = 0.01):
    onnx_model   = onnx.load(onnx_model_path)
    inits  = onnx_model.graph.initializer
    weights_map = {}
    for init in inits:
        W = numpy_helper.to_array(init)
        weights_map[init.name] = W
    print(weights_map.keys())
    for idx, key in enumerate(gradient_map.keys()):
        W_new  = weights_map[key] + learning_rate * gradient_map[key]
        tensor = numpy_helper.from_array(W_new)
        onnx_model.graph.initializer[idx].CopyFrom(tensor)
    onnx.save(onnx_model, onnx_updated_model_path)

[buddhapuneeth]

@jcwchen
I am able to resolve the issue by modifying to from_array as:
tensor = numpy_helper.from_array(W_new, key)

One more issue I faced was, new model was twice the old model because all weights changed from float 32 to float 64. This is due to 'gradient_map' in my code is float64, after correcting it worked fine.