onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/__init__.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""onnx2fx: Convert ONNX models to PyTorch FX GraphModules.
+
+This library provides tools for converting ONNX models into PyTorch FX
+GraphModules, enabling seamless integration with PyTorch's ecosystem for
+optimization, analysis, and deployment.
+
+Core Functions
+--------------
+convert : Convert an ONNX model to a PyTorch FX GraphModule.
+make_trainable : Convert buffers to trainable parameters for training.
+
+Model Analysis
+--------------
+analyze_model : Analyze an ONNX model for operator support.
+AnalysisResult : Dataclass containing analysis results.
+
+Operator Registration
+---------------------
+register_op : Register a custom operator handler.
+unregister_op : Unregister an operator handler.
+is_supported : Check if an operator is supported.
+get_supported_ops : List supported operators for a domain.
+get_all_supported_ops : Get all supported operators across all domains.
+get_registered_domains : Get list of registered domains.
+
+Exceptions
+----------
+Onnx2FxError : Base exception for all onnx2fx errors.
+UnsupportedOpError : Raised when an operator is not supported.
+ConversionError : Raised when conversion fails.
+ValueNotFoundError : Raised when a value is not found in environment.
+
+Example
+-------
+>>> import onnx
+>>> from onnx2fx import convert
+>>> model = onnx.load("model.onnx")
+>>> fx_module = convert(model)
+>>> # Use fx_module like any PyTorch module
+>>> output = fx_module(input_tensor)
+"""
+
+from importlib.metadata import PackageNotFoundError, version as _pkg_version
+
+from .converter import convert
+from .exceptions import (
+    Onnx2FxError,
+    UnsupportedOpError,
+    ConversionError,
+    ValueNotFoundError,
+    UnsupportedDTypeError,
+    ExternalDataError,
+    InferenceOnlyError,
+)
+from .op_registry import (
+    register_op,
+    unregister_op,
+    get_supported_ops,
+    get_all_supported_ops,
+    get_registered_domains,
+    is_supported,
+)
+from .utils.analyze import analyze_model, AnalysisResult
+from .utils.training import make_trainable
+
+try:
+    __version__ = _pkg_version("onnx2fx")
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.0.0"
+
+__all__ = [
+    "__version__",
+    # Core API
+    "convert",
+    # Training utilities
+    "make_trainable",
+    # Model analysis
+    "analyze_model",
+    "AnalysisResult",
+    # Operator registration
+    "register_op",
+    "unregister_op",
+    "get_supported_ops",
+    "get_all_supported_ops",
+    "get_registered_domains",
+    "is_supported",
+    # Exceptions
+    "Onnx2FxError",
+    "UnsupportedOpError",
+    "ConversionError",
+    "ValueNotFoundError",
+    "UnsupportedDTypeError",
+    "ExternalDataError",
+    "InferenceOnlyError",
+]

onnx2fx/converter.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Main entry point for converting ONNX models to PyTorch FX.
+
+This module provides the primary `convert` function for transforming
+ONNX models into equivalent PyTorch FX GraphModules.
+"""
+
+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
+    -------
+    torch.fx.GraphModule
+        A PyTorch FX Graph module.
+    """
+
+    if isinstance(model, str):
+        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")
+
+    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()

onnx2fx/exceptions.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Custom exceptions for onnx2fx."""
+
+
+class Onnx2FxError(Exception):
+    """Base exception for all onnx2fx errors.
+
+    This is the parent class for all custom exceptions in onnx2fx.
+    Users can catch this exception to handle any onnx2fx-related error.
+    """
+
+    pass
+
+
+class UnsupportedOpError(Onnx2FxError):
+    """Raised when an ONNX operator is not supported.
+
+    This exception is raised during conversion when the converter
+    encounters an ONNX operator that has no registered handler.
+
+    Parameters
+    ----------
+    op_type : str
+        The unsupported ONNX operator type.
+    domain : str, optional
+        The ONNX domain of the operator.
+    opset_version : int, optional
+        The opset version being used.
+    """
+
+    def __init__(
+        self,
+        op_type: str,
+        domain: str = "",
+        opset_version: int | None = None,
+    ):
+        self.op_type = op_type
+        self.domain = domain
+        self.opset_version = opset_version
+
+        domain_str = f" (domain: {domain})" if domain else ""
+        version_str = f" at opset {opset_version}" if opset_version else ""
+        message = f"Unsupported ONNX operator: {op_type}{domain_str}{version_str}"
+        super().__init__(message)
+
+
+class ConversionError(Onnx2FxError):
+    """Raised when conversion fails due to an error in the conversion process.
+
+    This exception is raised when the conversion process encounters
+    an error that prevents successful completion.
+
+    Parameters
+    ----------
+    message : str
+        A description of the conversion error.
+    node_name : str, optional
+        The name of the ONNX node where the error occurred.
+    op_type : str, optional
+        The ONNX operator type where the error occurred.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        node_name: str | None = None,
+        op_type: str | None = None,
+    ):
+        self.node_name = node_name
+        self.op_type = op_type
+
+        context_parts = []
+        if node_name:
+            context_parts.append(f"node: {node_name}")
+        if op_type:
+            context_parts.append(f"op: {op_type}")
+
+        if context_parts:
+            context = f" [{', '.join(context_parts)}]"
+        else:
+            context = ""
+
+        full_message = f"Conversion failed{context}: {message}"
+        super().__init__(full_message)
+
+
+class ValueNotFoundError(Onnx2FxError):
+    """Raised when a value is not found in the environment.
+
+    This exception is raised when trying to access a tensor value
+    that has not been defined in the conversion environment.
+
+    Parameters
+    ----------
+    name : str
+        The name of the value that was not found.
+    available : list[str], optional
+        List of available value names for debugging.
+    """
+
+    def __init__(self, name: str, available: list[str] | None = None):
+        self.name = name
+        self.available = available
+
+        message = f"Value '{name}' not found in environment"
+        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)