External data callback and ir.save fixes (#85)

This pull request introduces a new feature to the ONNX IR library by
adding support for a callback function in several methods. The callback
allows users to debug or log information about tensors being saved to
external data files (or enable a progress bar). Additionally, minor
improvements to code organization and comments are included.


![image](https://github.com/user-attachments/assets/0ced1a65-5235-469d-8053-2fa93116986a)

---------

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

src/onnx_ir/_io.py

@@ -7,10 +7,11 @@
 __all__ = ["load", "save"]
 
 import os
+from typing import Callable
 
 import onnx
 
-from onnx_ir import _core, serde
+from onnx_ir import _core, _protocols, serde
 from onnx_ir import external_data as _external_data
 from onnx_ir._polyfill import zip
 
@@ -43,6 +44,8 @@ def save(
     format: str | None = None,
     external_data: str | os.PathLike | None = None,
     size_threshold_bytes: int = 256,
+    callback: Callable[[_protocols.TensorProtocol, _external_data.CallbackInfo], None]
+    | None = None,
 ) -> None:
     """Save an ONNX model to a file.
 
@@ -52,6 +55,30 @@ def save(
     to load the newly saved model, or provide a different external data path that
     is not currently referenced by any tensors in the model.
 
+    .. tip::
+
+        A simple progress bar can be implemented by passing a callback function as the following::
+
+            import onnx_ir as ir
+            import tqdm
+
+            with tqdm.tqdm() as pbar:
+                total_set = False
+
+                def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) -> None:
+                    nonlocal total_set
+                    if not total_set:
+                        pbar.total = metadata.total
+                        total_set = True
+
+                    pbar.update()
+                    pbar.set_description(f"Saving {tensor.name} ({tensor.dtype}, {tensor.shape}) at offset {metadata.offset}")
+
+                ir.save(
+                    ...,
+                    callback=callback,
+                )
+
     Args:
         model: The model to save.
         path: The path to save the model to. E.g. "model.onnx".
@@ -65,6 +92,8 @@ def save(
             it will be serialized in the ONNX Proto message.
         size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
             Effective only when ``external_data`` is set.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
 
     Raises:
         ValueError: If the external data path is an absolute path.
@@ -77,12 +106,19 @@ def save(
         base_dir = os.path.dirname(path)
 
         # Store the original initializer values so they can be restored if modify_model=False
-        initializer_values = tuple(model.graph.initializers.values())
+        initializer_values: list[_core.Value] = []
+        for graph in model.graphs():
+            # Collect from all subgraphs as well
+            initializer_values.extend(graph.initializers.values())
         tensors = [v.const_value for v in initializer_values]
 
         try:
             model = _external_data.unload_from_model(
-                model, base_dir, external_data, size_threshold_bytes=size_threshold_bytes
+                model,
+                base_dir,
+                external_data,
+                size_threshold_bytes=size_threshold_bytes,
+                callback=callback,
             )
             proto = serde.serialize_model(model)
             onnx.save(proto, path, format=format)

src/onnx_ir/external_data.py

@@ -4,12 +4,15 @@
 
 from __future__ import annotations
 
+from typing import Callable
+
 __all__ = [
     "set_base_dir",
     "unload_from_model",
     "load_to_model",
     "convert_tensors_to_external",
     "convert_tensors_from_external",
+    "CallbackInfo",
 ]
 
 import dataclasses
@@ -48,6 +51,21 @@ class _ExternalDataInfo:
     length: int
 
 
+@dataclasses.dataclass
+class CallbackInfo:
+    """A class that shares information about a tensor that is to be saved as external data for callback functions.
+
+    Attributes:
+        total: The total number of tensors to save.
+        index: The index of the tensor being saved.
+        offset: The offset of the tensor in the external data file.
+    """
+
+    total: int
+    index: int
+    offset: int
+
+
 def _all_tensors(
     graph: _core.Graph | _core.GraphView, include_attributes: bool = False
 ) -> Iterator[_protocols.TensorProtocol]:
@@ -157,19 +175,34 @@ def _write_external_data(
     tensors: Sequence[_protocols.TensorProtocol],
     external_data_infos: Sequence[_ExternalDataInfo],
     file_path: str | os.PathLike,
+    callback: Callable[[_protocols.TensorProtocol, CallbackInfo], None] | None = None,
 ) -> None:
     """Write tensor data to an external file according to information stored in ExternalDataInfo objects.
 
     Args:
         tensors: Tensors to be written as external data.
         external_data_infos: External data information stored for each tensor to be written as external data.
         file_path: Location to which external data is to be stored.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
     """
-    assert len(tensors) == len(external_data_infos), (
+    tensors_count = len(tensors)
+    assert tensors_count == len(external_data_infos), (
         "Number of tensors and external data infos should match"
     )
     with open(file_path, "wb") as data_file:
-        for tensor, tensor_info in zip(tensors, external_data_infos, strict=True):
+        for i, (tensor, tensor_info) in enumerate(
+            zip(tensors, external_data_infos, strict=True)
+        ):
+            if callback is not None:
+                callback(
+                    tensor,
+                    CallbackInfo(
+                        total=tensors_count,
+                        index=i,
+                        offset=tensor_info.offset,
+                    ),
+                )
             current_offset = tensor_info.offset
             assert tensor is not None
             raw_data = tensor.tobytes()
@@ -228,6 +261,7 @@ def convert_tensors_to_external(
     tensors: Sequence[_protocols.TensorProtocol],
     base_dir: str | os.PathLike,
     relative_path: str | os.PathLike,
+    callback: Callable[[_protocols.TensorProtocol, CallbackInfo], None] | None = None,
 ) -> list[_core.ExternalTensor]:
     """Convert a sequence of any TensorProtocol tensors to external tensors.
 
@@ -238,6 +272,8 @@ def convert_tensors_to_external(
         tensors: Tensors to be converted to external tensors. They can be external tensors themselves.
         base_dir: Path of base directory.
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
 
     Returns:
         A list of external tensors derived from a list of input tensors. The order
@@ -285,7 +321,7 @@ def convert_tensors_to_external(
         external_info = _compute_external_data_info(tensor, current_offset)
         external_data_infos.append(external_info)
         current_offset = external_info.offset + external_info.length
-    _write_external_data(sorted_tensors, external_data_infos, path)
+    _write_external_data(sorted_tensors, external_data_infos, path, callback=callback)
 
     # Create external tensor objects
     external_tensors: list[_core.ExternalTensor] = [
@@ -336,6 +372,7 @@ def unload_from_model(
     relative_path: str | os.PathLike,
     *,
     size_threshold_bytes: int = 0,
+    callback: Callable[[_protocols.TensorProtocol, CallbackInfo], None] | None = None,
 ) -> _core.Model:
     """Convert all initializers equal or above size_threshold_bytes to external tensors in-place and save data to a single data file.
 
@@ -356,6 +393,8 @@ def unload_from_model(
         relative_path: Path to which external data is to be stored, relative to the ONNX file.
             E.g. "model.data"
         size_threshold_bytes: Save to external data if the tensor size in bytes is larger than this threshold.
+        callback: A callback function that is called for each tensor that is saved to external data
+            for debugging or logging purposes.
 
     Returns:
         An ir.Model with all initializer data equal or above ``size_threshold_bytes``
@@ -384,6 +423,7 @@ def unload_from_model(
         [v.const_value for v in initializers_to_become_external],  # type: ignore[misc]
         base_dir=base_dir,
         relative_path=relative_path,
+        callback=callback,
     )
 
     # Replace the initializer values with external tensors and save the model