torch-rechub 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

torch_rechub/utils/onnx_export.py

@@ -0,0 +1,220 @@
+"""
+ONNX Export Utilities for Torch-RecHub models.
+
+This module provides non-invasive ONNX export functionality for recommendation models.
+It uses reflection to extract feature information from models and wraps dict-input models
+to be compatible with ONNX's positional argument requirements.
+
+Date: 2024
+References:
+    - PyTorch ONNX Export: https://pytorch.org/docs/stable/onnx.html
+    - ONNX Runtime: https://onnxruntime.ai/docs/
+Authors: Torch-RecHub Contributors
+"""
+
+import os
+import warnings
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ..basic.features import DenseFeature, SequenceFeature, SparseFeature
+
+
+class ONNXWrapper(nn.Module):
+    """Wraps a dict-input model to accept positional arguments for ONNX compatibility.
+
+    ONNX does not support dict as input, so this wrapper converts positional arguments
+    back to dict format before passing to the original model.
+
+    Args:
+        model: The original PyTorch model that accepts dict input.
+        input_names: Ordered list of feature names corresponding to input positions.
+        mode: Optional mode for dual-tower models ("user" or "item").
+
+    Example:
+        >>> wrapper = ONNXWrapper(dssm_model, ["user_id", "movie_id", "hist_movie_id"])
+        >>> # Now can call: wrapper(user_id_tensor, movie_id_tensor, hist_tensor)
+    """
+
+    def __init__(self, model: nn.Module, input_names: List[str], mode: Optional[str] = None):
+        super().__init__()
+        self.model = model
+        self.input_names = input_names
+        self._original_mode = getattr(model, 'mode', None)
+
+        # Set mode for dual-tower models
+        if mode is not None and hasattr(model, 'mode'):
+            model.mode = mode
+
+    def forward(self, *args) -> torch.Tensor:
+        """Convert positional args to dict and call original model."""
+        if len(args) != len(self.input_names):
+            raise ValueError(f"Expected {len(self.input_names)} inputs, got {len(args)}. "
+                             f"Expected names: {self.input_names}")
+        x_dict = {name: arg for name, arg in zip(self.input_names, args)}
+        return self.model(x_dict)
+
+    def restore_mode(self):
+        """Restore the original mode of the model."""
+        if hasattr(self.model, 'mode'):
+            self.model.mode = self._original_mode
+
+
+# Re-export from model_utils for backward compatibility
+# The actual implementations are now in model_utils.py
+from .model_utils import extract_feature_info, generate_dummy_input, generate_dummy_input_dict, generate_dynamic_axes
+
+
+class ONNXExporter:
+    """Main class for exporting Torch-RecHub models to ONNX format.
+
+    This exporter handles the complexity of converting dict-input models to ONNX
+    by automatically extracting feature information and wrapping the model.
+
+    Args:
+        model: The PyTorch recommendation model to export.
+        device: Device for export operations (default: 'cpu').
+
+    Example:
+        >>> exporter = ONNXExporter(deepfm_model)
+        >>> exporter.export("model.onnx")
+
+        >>> # For dual-tower models
+        >>> exporter = ONNXExporter(dssm_model)
+        >>> exporter.export("user_tower.onnx", mode="user")
+        >>> exporter.export("item_tower.onnx", mode="item")
+    """
+
+    def __init__(self, model: nn.Module, device: str = 'cpu'):
+        self.model = model
+        self.device = device
+        self.feature_info = extract_feature_info(model)
+
+    def export(
+        self,
+        output_path: str,
+        mode: Optional[str] = None,
+        dummy_input: Optional[Dict[str,
+                                   torch.Tensor]] = None,
+        batch_size: int = 2,
+        seq_length: int = 10,
+        opset_version: int = 14,
+        dynamic_batch: bool = True,
+        verbose: bool = False
+    ) -> bool:
+        """Export the model to ONNX format.
+
+        Args:
+            output_path: Path to save the ONNX model.
+            mode: For dual-tower models, specify "user" or "item" to export
+                  only that tower. None exports the full model.
+            dummy_input: Optional dict of example inputs. If not provided,
+                         dummy inputs will be generated automatically.
+            batch_size: Batch size for generated dummy input (default: 2).
+            seq_length: Sequence length for SequenceFeature (default: 10).
+            opset_version: ONNX opset version (default: 14).
+            dynamic_batch: Whether to enable dynamic batch size (default: True).
+            verbose: Whether to print export details (default: False).
+
+        Returns:
+            True if export succeeded, False otherwise.
+
+        Raises:
+            RuntimeError: If ONNX export fails.
+        """
+        self.model.eval()
+        self.model.to(self.device)
+
+        # Determine which features to use based on mode
+        if mode == "user":
+            features = self.feature_info['user_features']
+            if not features:
+                raise ValueError("No user features found in model for mode='user'")
+        elif mode == "item":
+            features = self.feature_info['item_features']
+            if not features:
+                raise ValueError("No item features found in model for mode='item'")
+        else:
+            features = self.feature_info['features']
+
+        input_names = [f.name for f in features]
+
+        # Create wrapped model
+        wrapper = ONNXWrapper(self.model, input_names, mode=mode)
+        wrapper.eval()
+
+        # Generate or use provided dummy input
+        if dummy_input is not None:
+            dummy_tuple = tuple(dummy_input[name].to(self.device) for name in input_names)
+        else:
+            dummy_tuple = generate_dummy_input(features, batch_size=batch_size, seq_length=seq_length, device=self.device)
+
+        # Configure dynamic axes
+        dynamic_axes = None
+        if dynamic_batch:
+            seq_feature_names = [f.name for f in features if isinstance(f, SequenceFeature)]
+            dynamic_axes = generate_dynamic_axes(input_names=input_names, output_names=["output"], seq_features=seq_feature_names)
+
+        # Ensure output directory exists
+        output_dir = os.path.dirname(output_path)
+        if output_dir and not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        try:
+            with torch.no_grad():
+                torch.onnx.export(
+                    wrapper,
+                    dummy_tuple,
+                    output_path,
+                    input_names=input_names,
+                    output_names=["output"],
+                    dynamic_axes=dynamic_axes,
+                    opset_version=opset_version,
+                    do_constant_folding=True,
+                    verbose=verbose,
+                    dynamo=False  # Use legacy exporter for dynamic_axes support
+                )
+
+            if verbose:
+                print(f"Successfully exported ONNX model to: {output_path}")
+                print(f"  Input names: {input_names}")
+                print(f"  Opset version: {opset_version}")
+                print(f"  Dynamic batch: {dynamic_batch}")
+
+            return True
+
+        except Exception as e:
+            warnings.warn(f"ONNX export failed: {str(e)}")
+            raise RuntimeError(f"Failed to export ONNX model: {str(e)}") from e
+        finally:
+            # Restore original mode
+            wrapper.restore_mode()
+
+    def get_input_info(self, mode: Optional[str] = None) -> Dict[str, Any]:
+        """Get information about model inputs.
+
+        Args:
+            mode: For dual-tower models, "user" or "item".
+
+        Returns:
+            Dict with input names, types, and shapes.
+        """
+        if mode == "user":
+            features = self.feature_info['user_features']
+        elif mode == "item":
+            features = self.feature_info['item_features']
+        else:
+            features = self.feature_info['features']
+
+        info = []
+        for f in features:
+            feat_info = {'name': f.name, 'type': type(f).__name__, 'embed_dim': f.embed_dim}
+            if hasattr(f, 'vocab_size'):
+                feat_info['vocab_size'] = f.vocab_size
+            if hasattr(f, 'pooling'):
+                feat_info['pooling'] = f.pooling
+            info.append(feat_info)
+
+        return {'mode': mode, 'inputs': info, 'input_names': [f.name for f in features]}

torch_rechub/utils/visualization.py

@@ -0,0 +1,271 @@
+"""
+Model Visualization Utilities for Torch-RecHub.
+
+This module provides model structure visualization using torchview library.
+Requires optional dependencies: pip install torch-rechub[visualization]
+
+Example:
+    >>> from torch_rechub.utils.visualization import visualize_model, display_graph
+    >>> graph = visualize_model(model, depth=4)
+    >>> display_graph(graph)  # Display in Jupyter Notebook
+
+    >>> # Save to file
+    >>> visualize_model(model, save_path="model_arch.pdf")
+"""
+
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+# Check for optional dependencies
+TORCHVIEW_AVAILABLE = False
+TORCHVIEW_SKIP_REASON = "torchview not installed"
+
+try:
+    from torchview import draw_graph
+    TORCHVIEW_AVAILABLE = True
+except ImportError as e:
+    TORCHVIEW_SKIP_REASON = f"torchview not available: {e}"
+
+
+def _is_jupyter_environment() -> bool:
+    """Check if running in Jupyter/IPython environment."""
+    try:
+        from IPython import get_ipython
+        shell = get_ipython()
+        if shell is None:
+            return False
+        # Check for Jupyter notebook or qtconsole
+        shell_class = shell.__class__.__name__
+        return shell_class in ('ZMQInteractiveShell', 'TerminalInteractiveShell')
+    except (ImportError, NameError):
+        return False
+
+
+def display_graph(graph: Any, format: str = 'png') -> Any:
+    """Display a torchview ComputationGraph in Jupyter Notebook.
+
+    This function provides a reliable way to display visualization graphs
+    in Jupyter environments, especially VSCode Jupyter.
+
+    Parameters
+    ----------
+    graph : ComputationGraph
+        A torchview ComputationGraph object returned by visualize_model().
+    format : str, default='png'
+        Output format, 'png' recommended for VSCode.
+
+    Returns
+    -------
+    graphviz.Digraph or None
+        The displayed graph object, or None if display fails.
+
+    Examples
+    --------
+    >>> graph = visualize_model(model, depth=4)
+    >>> display_graph(graph)  # Works in VSCode Jupyter
+    """
+    if not TORCHVIEW_AVAILABLE:
+        raise ImportError(f"Visualization requires torchview. {TORCHVIEW_SKIP_REASON}\n"
+                          "Install with: pip install torch-rechub[visualization]")
+
+    try:
+        import graphviz
+
+        # Set format for VSCode compatibility
+        graphviz.set_jupyter_format(format)
+    except ImportError:
+        pass
+
+    # Get the visual_graph (graphviz.Digraph object)
+    visual = graph.visual_graph
+
+    # Try to use IPython display for explicit rendering
+    try:
+        from IPython.display import display
+        display(visual)
+        return visual
+    except ImportError:
+        # Not in IPython/Jupyter environment, return the graph
+        return visual
+
+
+def visualize_model(
+    model: nn.Module,
+    input_data: Optional[Dict[str,
+                              torch.Tensor]] = None,
+    batch_size: int = 2,
+    seq_length: int = 10,
+    depth: int = 3,
+    show_shapes: bool = True,
+    expand_nested: bool = True,
+    save_path: Optional[str] = None,
+    graph_name: str = "model",
+    device: str = "cpu",
+    dpi: int = 300,
+    **kwargs
+) -> Any:
+    """Visualize a Torch-RecHub model's computation graph.
+
+    This function generates a visual representation of the model architecture,
+    showing layer connections, tensor shapes, and nested module structures.
+    It automatically extracts feature information from the model to generate
+    appropriate dummy inputs.
+
+    Parameters
+    ----------
+    model : nn.Module
+        PyTorch model to visualize. Should be a Torch-RecHub model
+        with feature attributes (e.g., DeepFM, DSSM, MMOE).
+    input_data : dict, optional
+        Dict of example inputs {feature_name: tensor}.
+        If None, inputs are auto-generated based on model features.
+    batch_size : int, default=2
+        Batch size for auto-generated inputs.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature inputs.
+    depth : int, default=3
+        Visualization depth - higher values show more detail.
+        Set to -1 to show all layers.
+    show_shapes : bool, default=True
+        Whether to display tensor shapes on edges.
+    expand_nested : bool, default=True
+        Whether to expand nested nn.Module with dashed borders.
+    save_path : str, optional
+        Path to save the graph image. Supports .pdf, .svg, .png formats.
+        If None, displays in Jupyter or opens system viewer.
+    graph_name : str, default="model"
+        Name for the computation graph.
+    device : str, default="cpu"
+        Device for model execution during tracing.
+    dpi : int, default=300
+        Resolution in dots per inch for output image.
+        Higher values produce sharper images suitable for papers.
+    **kwargs : dict
+        Additional arguments passed to torchview.draw_graph().
+
+    Returns
+    -------
+    ComputationGraph
+        A torchview ComputationGraph object.
+        - Use `.visual_graph` property to get the graphviz.Digraph
+        - Use `.resize_graph(scale=1.5)` to adjust graph size
+
+    Raises
+    ------
+    ImportError
+        If torchview or graphviz is not installed.
+    ValueError
+        If model has no recognizable feature attributes.
+
+    Notes
+    -----
+    Default Display Behavior:
+        When `save_path` is None (default):
+        - In Jupyter/IPython: automatically displays the graph inline
+        - In Python script: opens the graph with system default viewer
+
+    Requires graphviz system package: apt/brew/choco install graphviz.
+    For Jupyter display issues, try: graphviz.set_jupyter_format('png').
+
+    Examples
+    --------
+    >>> from torch_rechub.models.ranking import DeepFM
+    >>> from torch_rechub.utils.visualization import visualize_model
+    >>>
+    >>> # Auto-display in Jupyter or open in viewer
+    >>> visualize_model(model, depth=4)  # No save_path needed
+    >>>
+    >>> # Save to high-DPI PNG for paper
+    >>> visualize_model(model, save_path="model.png", dpi=300)
+    """
+    if not TORCHVIEW_AVAILABLE:
+        raise ImportError(
+            f"Visualization requires torchview. {TORCHVIEW_SKIP_REASON}\n"
+            "Install with: pip install torch-rechub[visualization]\n"
+            "Also ensure graphviz is installed on your system:\n"
+            "  - Ubuntu/Debian: sudo apt-get install graphviz\n"
+            "  - macOS: brew install graphviz\n"
+            "  - Windows: choco install graphviz"
+        )
+
+    # Import feature extraction utilities from model_utils
+    from .model_utils import extract_feature_info, generate_dummy_input_dict
+
+    model.eval()
+    model.to(device)
+
+    # Auto-generate input data if not provided
+    if input_data is None:
+        feature_info = extract_feature_info(model)
+        features = feature_info['features']
+
+        if not features:
+            raise ValueError("Could not extract feature information from model. "
+                             "Please provide input_data parameter manually.")
+
+        # Generate dummy input dict
+        input_data = generate_dummy_input_dict(features, batch_size=batch_size, seq_length=seq_length, device=device)
+    else:
+        # Ensure input tensors are on correct device
+        input_data = {k: v.to(device) for k, v in input_data.items()}
+
+    # IMPORTANT: Wrap input_data dict in a tuple to work around torchview's behavior
+    #
+    # torchview's forward_prop function checks the type of input_data:
+    #   - If isinstance(x, (list, tuple)): model(*x)
+    #   - If isinstance(x, Mapping): model(**x)  <- This unpacks dict as kwargs!
+    #   - Else: model(x)
+    #
+    # torch-rechub models expect forward(self, x) where x is a complete dict.
+    # By wrapping the dict in a tuple, torchview will call:
+    #   model(*(input_dict,)) = model(input_dict)
+    # which is exactly what our models expect.
+    input_data_wrapped = (input_data,)
+
+    # Call torchview.draw_graph without saving (we'll save manually with DPI)
+    graph = draw_graph(
+        model,
+        input_data=input_data_wrapped,
+        graph_name=graph_name,
+        depth=depth,
+        device=device,
+        expand_nested=expand_nested,
+        show_shapes=show_shapes,
+        save_graph=False,  # Don't save here, we'll save manually with DPI
+        **kwargs
+    )
+
+    # Set DPI for high-quality output (must be set BEFORE rendering/saving)
+    graph.visual_graph.graph_attr['dpi'] = str(dpi)
+
+    # Handle save_path: manually save with DPI applied
+    if save_path:
+        import os
+        directory = os.path.dirname(save_path) or "."
+        filename = os.path.splitext(os.path.basename(save_path))[0]
+        ext = os.path.splitext(save_path)[1].lstrip('.')
+        # Default to pdf if no extension
+        output_format = ext if ext else 'pdf'
+        # Create directory if it doesn't exist
+        if directory != "." and not os.path.exists(directory):
+            os.makedirs(directory, exist_ok=True)
+        # Render and save with DPI applied
+        graph.visual_graph.render(
+            filename=filename,
+            directory=directory,
+            format=output_format,
+            cleanup=True  # Remove intermediate .gv file
+        )
+
+    # Handle default display behavior when save_path is None
+    if save_path is None:
+        if _is_jupyter_environment():
+            # In Jupyter: display inline
+            display_graph(graph)
+        else:
+            # In script: open with system viewer
+            graph.visual_graph.view(cleanup=True)
+
+    return graph