large onnx file support

Author: mshr-h

src/onnx2fx/__init__.py

@@ -47,6 +47,9 @@
     UnsupportedOpError,
     ConversionError,
     ValueNotFoundError,
+    UnsupportedDTypeError,
+    ExternalDataError,
+    InferenceOnlyError,
 )
 from .op_registry import (
     register_op,
@@ -79,4 +82,7 @@
     "UnsupportedOpError",
     "ConversionError",
     "ValueNotFoundError",
+    "UnsupportedDTypeError",
+    "ExternalDataError",
+    "InferenceOnlyError",
 ]

src/onnx2fx/converter.py

@@ -5,23 +5,34 @@
 ONNX models into equivalent PyTorch FX GraphModules.
 """
 
-from typing import Union
+import os
+from typing import Optional, Union
 
 import onnx
 import torch
 
 from .graph_builder import GraphBuilder
+from .utils.external_data import validate_external_data_model
 
 
 def convert(
     model: Union[onnx.ModelProto, str],
+    *,
+    base_dir: Optional[str] = None,
+    memmap_external_data: bool = False,
 ) -> torch.fx.GraphModule:
     """Convert an ONNX model into a ``torch.fx.GraphModule``.
 
     Parameters
     ----------
     model : Union[onnx.ModelProto, str]
         Either an in-memory ``onnx.ModelProto`` or a file path to an ONNX model.
+    base_dir : Optional[str], optional
+        Base directory for resolving external data tensors. Required when
+        ``memmap_external_data=True`` and a relative external data path is used.
+    memmap_external_data : bool, optional
+        If True, do not load external data into memory. Instead, keep external
+        data references for memmap-based loading during conversion.
 
     Returns
     -------
@@ -30,10 +41,22 @@ def convert(
     """
 
     if isinstance(model, str):
-        model = onnx.load(model)
+        if base_dir is None:
+            base_dir = os.path.dirname(os.path.abspath(model))
+        if memmap_external_data:
+            model = onnx.load_model(model, load_external_data=False)
+        else:
+            model = onnx.load_model(model)
     elif isinstance(model, onnx.ModelProto):
         model = model
     else:
         raise TypeError("model must be a path or onnx.ModelProto")
 
-    return GraphBuilder(model).build()
+    if memmap_external_data:
+        validate_external_data_model(model, base_dir=base_dir, strict=True)
+
+    return GraphBuilder(
+        model,
+        base_dir=base_dir,
+        memmap_external_data=memmap_external_data,
+    ).build()

src/onnx2fx/exceptions.py

@@ -106,3 +106,50 @@ def __init__(self, name: str, available: list[str] | None = None):
         if available:
             message += f". Available: {available}"
         super().__init__(message)
+
+
+class UnsupportedDTypeError(Onnx2FxError):
+    """Raised when an ONNX tensor dtype is not supported.
+
+    Parameters
+    ----------
+    onnx_dtype : int
+        ONNX TensorProto data type enum value.
+    tensor_name : str
+        Name of the tensor.
+    details : str, optional
+        Additional details about the failure.
+    """
+
+    def __init__(self, onnx_dtype: int, tensor_name: str, details: str | None = None):
+        self.onnx_dtype = onnx_dtype
+        self.tensor_name = tensor_name
+        self.details = details
+
+        dtype_name = f"{onnx_dtype}"
+        try:
+            import onnx
+
+            dtype_name = onnx.TensorProto.DataType.Name(onnx_dtype)
+        except Exception:
+            pass
+
+        message = f"Unsupported dtype for tensor '{tensor_name}': {dtype_name}"
+        if details:
+            message += f" ({details})"
+        super().__init__(message)
+
+
+class ExternalDataError(Onnx2FxError):
+    """Raised when external data metadata is invalid or inaccessible."""
+
+    def __init__(self, tensor_name: str, message: str):
+        self.tensor_name = tensor_name
+        super().__init__(f"External data error for '{tensor_name}': {message}")
+
+
+class InferenceOnlyError(Onnx2FxError):
+    """Raised when an inference-only model is used for training."""
+
+    def __init__(self, message: str):
+        super().__init__(message)

src/onnx2fx/graph_builder.py

@@ -2,14 +2,17 @@
 from collections import deque
 from typing import Any, Callable, Dict, List, Optional, Tuple, Sequence, Union
 
+import numpy as np
+
 import torch
 import torch.fx
 import onnx
 from onnx import numpy_helper
 
-from .exceptions import UnsupportedOpError, ValueNotFoundError
+from .exceptions import UnsupportedDTypeError, UnsupportedOpError, ValueNotFoundError
 from .op_registry import get_handler
 from .utils.dtype import DTYPE_MAP
+from .utils.external_data import resolve_external_data
 from .utils.names import sanitize_name
 
 # Import ops module to register all operators
@@ -178,7 +181,13 @@ class GraphBuilder:
         The opset version for the default ONNX domain.
     """
 
-    def __init__(self, model: onnx.ModelProto) -> None:
+    def __init__(
+        self,
+        model: onnx.ModelProto,
+        *,
+        base_dir: Optional[str] = None,
+        memmap_external_data: bool = False,
+    ) -> None:
         # Try shape inference but preserve original model if it fails
         # (shape_inference may drop graph contents for large models with external data)
         try:
@@ -191,6 +200,8 @@ def __init__(self, model: onnx.ModelProto) -> None:
             pass
         self.model: onnx.ModelProto = model
         self.graph: torch.fx.Graph = torch.fx.Graph()
+        self._base_dir = base_dir
+        self._memmap_external_data = memmap_external_data
         self.value_info_map = self._create_value_info_map()
         self.initializer_map = self._create_initializer_map()
         self.input_names: List[str] = []
@@ -299,6 +310,8 @@ def build(self) -> torch.fx.GraphModule:
         for name, submod in self._submodules.items():
             root_module.add_module(name, submod)
         module = torch.fx.GraphModule(root_module, self.graph)
+        if self._memmap_external_data:
+            module._onnx2fx_inference_only = True
         module.graph.lint()
         return module
 
@@ -456,7 +469,19 @@ def extract_tensor_shape(
         def extract_tensor_dtype(value: onnx.ValueInfoProto) -> Optional[torch.dtype]:
             """Extract the Torch dtype that corresponds to a value info."""
 
-            return DTYPE_MAP.get(value.type.tensor_type.elem_type)
+            onnx_dtype = value.type.tensor_type.elem_type
+            if onnx_dtype == 0:
+                return None
+            torch_dtype = DTYPE_MAP.get(onnx_dtype)
+            if torch_dtype is None:
+                if onnx_dtype == onnx.TensorProto.STRING:
+                    return None
+                raise UnsupportedDTypeError(
+                    onnx_dtype=onnx_dtype,
+                    tensor_name=value.name,
+                    details="value_info dtype not supported",
+                )
+            return torch_dtype
 
         info_map = {}
         for value_info in (
@@ -501,10 +526,39 @@ def _create_initializer_map(self) -> Dict[str, torch.Tensor]:
         """Build a mapping from initializer names to PyTorch tensors."""
         init_map = {}
         for initializer in self.model.graph.initializer:
-            np_array = numpy_helper.to_array(initializer)
-            init_map[initializer.name] = torch.from_numpy(np_array.copy())
+            init_map[initializer.name] = self.load_tensor(initializer)
         return init_map
 
+    def load_tensor(self, tensor: onnx.TensorProto) -> torch.Tensor:
+        """Load an ONNX TensorProto into a Torch tensor."""
+        onnx_dtype = tensor.data_type
+        if DTYPE_MAP.get(onnx_dtype) is None:
+            raise UnsupportedDTypeError(
+                onnx_dtype=onnx_dtype,
+                tensor_name=tensor.name or "<unnamed>",
+                details="initializer dtype not supported",
+            )
+
+        if self._memmap_external_data and (
+            tensor.data_location == onnx.TensorProto.EXTERNAL or tensor.external_data
+        ):
+            info = resolve_external_data(
+                tensor,
+                base_dir=self._base_dir,
+                strict=True,
+            )
+            memmap_array = np.memmap(
+                info.path,
+                dtype=info.numpy_dtype,
+                mode="r",
+                offset=info.offset,
+                shape=info.shape,
+            )
+            return torch.from_numpy(memmap_array)
+
+        np_array = numpy_helper.to_array(tensor)
+        return torch.from_numpy(np_array.copy())
+
     def _load_initializers(self) -> None:
         """Load ONNX initializers as constant nodes in the FX graph."""
         for name, tensor in self.initializer_map.items():

src/onnx2fx/ops/control_flow.py

@@ -4,7 +4,7 @@
 This module implements ONNX control flow operators like Loop and If.
 """
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
 
 import onnx
 import torch
@@ -70,6 +70,7 @@ def _build_subgraph_module(
     parent_env: Dict[str, torch.fx.Node],
     parent_opset_versions: Dict[str, int],
     parent_type_info: Optional[Dict[str, bool]] = None,
+    tensor_loader: Optional[Callable[[onnx.TensorProto], torch.Tensor]] = None,
 ) -> Tuple[torch.fx.GraphModule, List[str], List[str], List[str]]:
     """Build an FX GraphModule from an ONNX subgraph.
 
@@ -106,8 +107,11 @@ def _build_subgraph_module(
     # Load initializers from subgraph
     initializer_map: Dict[str, torch.Tensor] = {}
     for initializer in body_graph.initializer:
-        np_array = numpy_helper.to_array(initializer)
-        initializer_map[initializer.name] = torch.from_numpy(np_array.copy())
+        if tensor_loader is not None:
+            initializer_map[initializer.name] = tensor_loader(initializer)
+        else:
+            np_array = numpy_helper.to_array(initializer)
+            initializer_map[initializer.name] = torch.from_numpy(np_array.copy())
 
     # Register initializers as constants
     for name, tensor in initializer_map.items():
@@ -159,9 +163,16 @@ def __init__(self):
             self.initializer_map = initializer_map
             self._body_graph = body_graph
             self._parent_type_info = parent_type_info
+            self._tensor_loader = tensor_loader
             # Build type info for this subgraph (to pass to nested subgraphs)
             self._type_info = self._build_type_info()
 
+        def load_tensor(self, tensor: onnx.TensorProto) -> torch.Tensor:
+            if self._tensor_loader is not None:
+                return self._tensor_loader(tensor)
+            np_array = numpy_helper.to_array(tensor)
+            return torch.from_numpy(np_array.copy())
+
         def _build_type_info(self) -> Dict[str, bool]:
             """Build a mapping of value names to whether they are optional types."""
             info: Dict[str, bool] = {}
@@ -437,7 +448,11 @@ def loop_op(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
     # Build subgraph module
     body_module, body_input_names, body_output_names, outer_refs = (
         _build_subgraph_module(
-            body_graph, builder.env, builder._opset_versions, parent_type_info
+            body_graph,
+            builder.env,
+            builder._opset_versions,
+            parent_type_info,
+            tensor_loader=builder.load_tensor,
         )
     )
 
@@ -628,7 +643,11 @@ def scan_op(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
     # Build subgraph module
     body_module, body_input_names, body_output_names, outer_refs = (
         _build_subgraph_module(
-            body_graph, builder.env, builder._opset_versions, parent_type_info
+            body_graph,
+            builder.env,
+            builder._opset_versions,
+            parent_type_info,
+            tensor_loader=builder.load_tensor,
         )
     )
 
@@ -712,7 +731,11 @@ def scan_op_v8(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
     # Build subgraph module
     body_module, body_input_names, body_output_names, outer_refs = (
         _build_subgraph_module(
-            body_graph, builder.env, builder._opset_versions, parent_type_info
+            body_graph,
+            builder.env,
+            builder._opset_versions,
+            parent_type_info,
+            tensor_loader=builder.load_tensor,
         )
     )
 
@@ -884,12 +907,20 @@ def if_op(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
     # Build subgraph modules for both branches
     then_module, then_input_names, then_output_names, then_outer_refs = (
         _build_subgraph_module(
-            then_graph, builder.env, builder._opset_versions, parent_type_info
+            then_graph,
+            builder.env,
+            builder._opset_versions,
+            parent_type_info,
+            tensor_loader=builder.load_tensor,
         )
     )
     else_module, else_input_names, else_output_names, else_outer_refs = (
         _build_subgraph_module(
-            else_graph, builder.env, builder._opset_versions, parent_type_info
+            else_graph,
+            builder.env,
+            builder._opset_versions,
+            parent_type_info,
+            tensor_loader=builder.load_tensor,
         )
     )
 

src/onnx2fx/ops/tensor.py

@@ -24,7 +24,7 @@
 @register("Constant")
 def constant(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
     """Create a constant tensor."""
-    value = get_attribute(node, "value")
+    value = get_attribute(node, "value", tensor_loader=builder.load_tensor)
     if value is None:
         value_float = get_attribute(node, "value_float")
         if value_float is not None:
@@ -781,7 +781,7 @@ def size(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
 def constant_of_shape(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
     """Create tensor filled with constant value."""
     shape = builder.get_value(node.input[0])
-    value = get_attribute(node, "value")
+    value = get_attribute(node, "value", tensor_loader=builder.load_tensor)
 
     if value is not None:
         fill_value = (