.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_core/plot_raw_data_callback.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_core_plot_raw_data_callback.py: .. _l-example-plot-raw-data-callback: Track tensor weights while parsing with a raw_data callback =========================================================== This example shows how to use :attr:`onnx_light.onnx.ParseOptions.raw_data_callback` to hook into model parsing. The callback is invoked for every :class:`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. .. GENERATED FROM PYTHON SOURCE LINES 20-27 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 28-34 Build a tiny model with a couple of weight tensors -------------------------------------------------- Both initializers are stored as ``raw_data`` (this is what :func:`onnx_light.onnx.numpy_helper.from_array` produces), so the callback fires once per tensor. .. GENERATED FROM PYTHON SOURCE LINES 34-45 .. code-block:: Python 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)}") .. rst-class:: sphx-glr-script-out .. code-block:: none Number of initializers: 2 .. GENERATED FROM PYTHON SOURCE LINES 46-52 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. .. GENERATED FROM PYTHON SOURCE LINES 52-67 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 68-74 Parse the model with the callback enabled ----------------------------------------- The callback is set on :class:`onnx_light.onnx.ParseOptions` and passed to ``ParseFromString``. Because nested protos share the same options, it fires for every tensor in the whole model. .. GENERATED FROM PYTHON SOURCE LINES 74-84 .. code-block:: Python 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") .. rst-class:: sphx-glr-script-out .. code-block:: none parsed tensor 'W0': 256 bytes of raw_data parsed tensor 'W1': 32 bytes of raw_data .. GENERATED FROM PYTHON SOURCE LINES 85-86 The tensor data is fully usable: attaching a deleter does not move the bytes. .. GENERATED FROM PYTHON SOURCE LINES 86-91 .. code-block:: Python 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.") .. rst-class:: sphx-glr-script-out .. code-block:: none Tensor values match the originals. .. GENERATED FROM PYTHON SOURCE LINES 92-97 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. .. GENERATED FROM PYTHON SOURCE LINES 97-106 .. code-block:: Python del parsed_model import gc gc.collect() print(f"released tensors: {released_tensors}") .. rst-class:: sphx-glr-script-out .. code-block:: none released tensors: ['W0', 'W1'] .. GENERATED FROM PYTHON SOURCE LINES 107-112 Read-only inspection -------------------- Returning ``None`` from the callback leaves ownership untouched, so it can be used purely to inspect tensors as they are parsed. .. GENERATED FROM PYTHON SOURCE LINES 112-130 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none inspected tensor names: ['W0', 'W1'] .. GENERATED FROM PYTHON SOURCE LINES 131-139 Reporting progress with ``RawDataCallback`` ------------------------------------------- :class:`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. .. GENERATED FROM PYTHON SOURCE LINES 139-151 .. code-block:: Python 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)) .. rst-class:: sphx-glr-script-out .. code-block:: none W0: 256 bytes W1: 32 bytes .. GENERATED FROM PYTHON SOURCE LINES 152-164 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 :attr:`ParseOptions.raw_data_callback`. .. note:: This requires an onnx-light build with OpenSSL support (the same requirement as :func:`onnx_light.onnx.save_encrypted_string`). .. GENERATED FROM PYTHON SOURCE LINES 164-232 .. code-block:: Python 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.") .. rst-class:: sphx-glr-script-out .. code-block:: none ChaCha20 callback encryption/decryption round-trip succeeded. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.187 seconds) .. _sphx_glr_download_auto_examples_core_plot_raw_data_callback.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_raw_data_callback.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_raw_data_callback.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_raw_data_callback.zip ` .. include:: plot_raw_data_callback.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_