Note
Go to the end to download the full example code.
Track tensor weights while parsing with a raw_data callback#
This example shows how to use onnx_light.onnx.ParseOptions.raw_data_callback
to hook into model parsing. The callback is invoked for every
onnx_light.onnx.TensorProto once its raw_data has been resolved
(including external-data tensors, after their bytes have been loaded).
For each tensor the callback may return a deleter — a zero-argument callable
that runs when the tensor’s raw_data is released (when the model and every
copy sharing the same buffer go out of scope). This lets callers take custom
ownership of tensor data and register matching cleanup, regardless of whether
the bytes live on disk (no_copy view of an mmap or external file) or in CPU
memory. Returning None leaves the tensor’s ownership unchanged, which makes
the callback equally usable as a read-only inspection hook.
import numpy as np
import onnx_light.onnx.helper as oh
import onnx_light.onnx.numpy_helper as onh
import onnx_light.onnx as onnxl
Build a tiny model with a couple of weight tensors#
Both initializers are stored as raw_data (this is what
onnx_light.onnx.numpy_helper.from_array() produces), so the callback
fires once per tensor.
w0 = np.random.randn(8, 8).astype(np.float32)
w1 = np.random.randn(8).astype(np.float32)
initializers = [onh.from_array(w0, name="W0"), onh.from_array(w1, name="W1")]
graph = oh.make_graph([], "demo_graph", [], [], initializer=initializers)
onnx_model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)], ir_version=9)
serialized = onnx_model.SerializeToString()
print(f"Number of initializers: {len(onnx_model.graph.initializer)}")
Number of initializers: 2
Define the parsing callback#
raw_data_callback receives the freshly parsed tensor. Here we record the
name and byte size of every tensor and return a deleter that notes when the
underlying buffer is released.
parsed_tensors = []
released_tensors = []
def on_raw_data(tensor: onnxl.TensorProto):
"""Records the tensor and returns a deleter run when its raw_data is freed."""
parsed_tensors.append((tensor.name, len(tensor.raw_data)))
def deleter():
released_tensors.append(tensor.name)
return deleter
Parse the model with the callback enabled#
The callback is set on onnx_light.onnx.ParseOptions and passed to
ParseFromString. Because nested protos share the same options, it fires
for every tensor in the whole model.
options = onnxl.ParseOptions()
options.raw_data_callback = on_raw_data
parsed_model = onnxl.ModelProto()
parsed_model.ParseFromString(serialized, options)
for name, size in parsed_tensors:
print(f"parsed tensor {name!r}: {size} bytes of raw_data")
parsed tensor 'W0': 256 bytes of raw_data
parsed tensor 'W1': 32 bytes of raw_data
The tensor data is fully usable: attaching a deleter does not move the bytes.
np.testing.assert_array_equal(onh.to_array(parsed_model.graph.initializer[0]), w0)
np.testing.assert_array_equal(onh.to_array(parsed_model.graph.initializer[1]), w1)
print("Tensor values match the originals.")
Tensor values match the originals.
Releasing the model triggers the registered deleters#
When the parsed model (and every tensor sharing the same buffer) is released, each deleter runs exactly once.
del parsed_model
import gc
gc.collect()
print(f"released tensors: {released_tensors}")
released tensors: ['W0', 'W1']
Read-only inspection#
Returning None from the callback leaves ownership untouched, so it can be
used purely to inspect tensors as they are parsed.
names = []
inspect_options = onnxl.ParseOptions()
def record_name(tensor: onnxl.TensorProto):
"""Records the tensor name and returns None to keep ownership unchanged."""
names.append(tensor.name)
return None
inspect_options.raw_data_callback = record_name
inspected = onnxl.ModelProto()
inspected.ParseFromString(serialized, inspect_options)
print(f"inspected tensor names: {names}")
inspected tensor names: ['W0', 'W1']
Reporting progress with RawDataCallback#
onnx_light.onnx.RawDataCallback is a ready-made callback object for
the common case of only reporting progress: it forwards every parsed tensor
to an on_tensor callable and always returns None, so the default C++
allocation is left untouched. Assign an instance to raw_data_callback to
print progress without changing how tensor data is owned.
progress = []
progress_options = onnxl.ParseOptions()
progress_options.raw_data_callback = onnxl.RawDataCallback(
lambda tensor: progress.append(f"{tensor.name}: {len(tensor.raw_data)} bytes")
)
with_progress = onnxl.ModelProto()
with_progress.ParseFromString(serialized, progress_options)
print("\n".join(progress))
W0: 256 bytes
W1: 32 bytes
Encrypting weights with a serialization callback (ChaCha20) and restoring them on parse#
SerializeOptions.raw_data_callback can also rewrite the raw bytes. The
snippet below encrypts every initializer payload with ONNXCRY2
(ChaCha20-Poly1305) during serialization, stores the encrypted blob in
raw_data (as UINT8), then restores the original tensor bytes during
parsing with ParseOptions.raw_data_callback.
Note
This requires an onnx-light build with OpenSSL support (the same
requirement as onnx_light.onnx.save_encrypted_string()).
if hasattr(onnxl.ModelProto(), "SerializeToEncryptedString"):
secret = "callback-secret"
def _encrypt_bytes_chacha20(raw: bytes) -> bytes:
"""Encrypts raw bytes into an ONNXCRY2 blob."""
payload = onh.from_array(np.frombuffer(raw, dtype=np.uint8), name="PAYLOAD")
holder = oh.make_model(oh.make_graph([], "encrypt_payload", [], [], [payload]))
return onnxl.save_encrypted_string(holder, secret, encryption="ChaCha20-Poly1305")
def _decrypt_bytes_chacha20(blob: bytes) -> bytes:
"""Decrypts an ONNXCRY2 blob back to raw bytes."""
holder = onnxl.load_encrypted_string(blob, secret)
return onh.to_array(holder.graph.initializer[0]).tobytes()
encrypted_by_name = {}
meta_by_name = {}
def encrypt_weights(tensor: onnxl.TensorProto, buffer, size_only: bool) -> int:
"""Encrypts tensor bytes during callback serialization.
Returns:
The number of bytes written (or expected when size_only is True).
"""
encrypted = encrypted_by_name.get(tensor.name)
if encrypted is None:
encrypted = _encrypt_bytes_chacha20(bytes(tensor.raw_data))
encrypted_by_name[tensor.name] = encrypted
meta_by_name[tensor.name] = (int(tensor.data_type), list(tensor.dims))
if size_only:
return len(encrypted)
orig_dtype, orig_dims = meta_by_name[tensor.name]
tensor.data_type = onnxl.TensorProto.UINT8
tensor.ClearField("dims")
tensor.dims.extend([len(encrypted)])
tensor.doc_string = f"chacha20:{orig_dtype}:{','.join(map(str, orig_dims))}"
np.copyto(np.asarray(buffer), np.frombuffer(encrypted, dtype=np.uint8))
return len(encrypted)
serialize_options = onnxl.SerializeOptions()
serialize_options.raw_data_callback = encrypt_weights
encrypted_serialized = onnx_model.SerializeToString(serialize_options)
def decrypt_weights(tensor: onnxl.TensorProto) -> None:
"""Restores original tensor metadata and raw_data during parsing.
Returns:
None to keep the default ownership behavior.
"""
if not tensor.doc_string.startswith("chacha20:"):
return None
_, dtype_text, dims_text = tensor.doc_string.split(":", 2)
restored = _decrypt_bytes_chacha20(bytes(tensor.raw_data))
tensor.data_type = int(dtype_text)
tensor.ClearField("dims")
if dims_text:
tensor.dims.extend(int(value) for value in dims_text.split(",") if value)
tensor.raw_data = restored
return None
parse_options = onnxl.ParseOptions()
parse_options.raw_data_callback = decrypt_weights
decrypted_model = onnxl.ModelProto()
decrypted_model.ParseFromString(encrypted_serialized, parse_options)
np.testing.assert_array_equal(onh.to_array(decrypted_model.graph.initializer[0]), w0)
np.testing.assert_array_equal(onh.to_array(decrypted_model.graph.initializer[1]), w1)
print("ChaCha20 callback encryption/decryption round-trip succeeded.")
ChaCha20 callback encryption/decryption round-trip succeeded.
Total running time of the script: (0 minutes 0.187 seconds)
Related examples
Measures loading and saving time for an ONNX model
Benchmark streaming vs in-memory alignment of external data