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

torch_rechub/trainers/ctr_trainer.py

@@ -189,3 +189,100 @@ class CTRTrainer(object):
 
         exporter = ONNXExporter(model, device=export_device)
         return exporter.export(output_path=output_path, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+
+    def visualization(self, input_data=None, batch_size=2, seq_length=10, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
+        """Visualize the model's computation graph.
+
+        This method 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.
+
+        Parameters
+        ----------
+        input_data : dict, optional
+            Example input dict {feature_name: tensor}.
+            If not provided, dummy inputs will be generated automatically.
+        batch_size : int, default=2
+            Batch size for auto-generated dummy input.
+        seq_length : int, default=10
+            Sequence length for SequenceFeature.
+        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.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        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.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+
+        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
+
+        Examples
+        --------
+        >>> trainer = CTRTrainer(model, ...)
+        >>> trainer.fit(train_dl, val_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        from ..utils.visualization import TORCHVIEW_AVAILABLE, visualize_model
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "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"
+            )
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        return visualize_model(
+            model,
+            input_data=input_data,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            depth=depth,
+            show_shapes=show_shapes,
+            expand_nested=expand_nested,
+            save_path=save_path,
+            graph_name=graph_name,
+            device=viz_device,
+            dpi=dpi,
+            **kwargs
+        )

torch_rechub/trainers/match_trainer.py

@@ -237,3 +237,100 @@ class MatchTrainer(object):
             # Restore original mode
             if hasattr(model, 'mode'):
                 model.mode = original_mode
+
+    def visualization(self, input_data=None, batch_size=2, seq_length=10, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
+        """Visualize the model's computation graph.
+
+        This method 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.
+
+        Parameters
+        ----------
+        input_data : dict, optional
+            Example input dict {feature_name: tensor}.
+            If not provided, dummy inputs will be generated automatically.
+        batch_size : int, default=2
+            Batch size for auto-generated dummy input.
+        seq_length : int, default=10
+            Sequence length for SequenceFeature.
+        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.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        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.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+
+        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
+
+        Examples
+        --------
+        >>> trainer = MatchTrainer(model, ...)
+        >>> trainer.fit(train_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        from ..utils.visualization import TORCHVIEW_AVAILABLE, visualize_model
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "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"
+            )
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        return visualize_model(
+            model,
+            input_data=input_data,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            depth=depth,
+            show_shapes=show_shapes,
+            expand_nested=expand_nested,
+            save_path=save_path,
+            graph_name=graph_name,
+            device=viz_device,
+            dpi=dpi,
+            **kwargs
+        )

torch_rechub/trainers/mtl_trainer.py

@@ -257,3 +257,100 @@ class MTLTrainer(object):
 
         exporter = ONNXExporter(model, device=export_device)
         return exporter.export(output_path=output_path, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+
+    def visualization(self, input_data=None, batch_size=2, seq_length=10, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
+        """Visualize the model's computation graph.
+
+        This method 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.
+
+        Parameters
+        ----------
+        input_data : dict, optional
+            Example input dict {feature_name: tensor}.
+            If not provided, dummy inputs will be generated automatically.
+        batch_size : int, default=2
+            Batch size for auto-generated dummy input.
+        seq_length : int, default=10
+            Sequence length for SequenceFeature.
+        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.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        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.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+
+        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
+
+        Examples
+        --------
+        >>> trainer = MTLTrainer(model, task_types=["classification", "classification"])
+        >>> trainer.fit(train_dl, val_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        from ..utils.visualization import TORCHVIEW_AVAILABLE, visualize_model
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "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"
+            )
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        return visualize_model(
+            model,
+            input_data=input_data,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            depth=depth,
+            show_shapes=show_shapes,
+            expand_nested=expand_nested,
+            save_path=save_path,
+            graph_name=graph_name,
+            device=viz_device,
+            dpi=dpi,
+            **kwargs
+        )

torch_rechub/trainers/seq_trainer.py

@@ -291,3 +291,137 @@ class SeqTrainer(object):
         except Exception as e:
             warnings.warn(f"ONNX export failed: {str(e)}")
             raise RuntimeError(f"Failed to export ONNX model: {str(e)}") from e
+
+    def visualization(self, seq_length=50, vocab_size=None, batch_size=2, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
+        """Visualize the model's computation graph.
+
+        This method generates a visual representation of the sequence model
+        architecture, showing layer connections, tensor shapes, and nested
+        module structures.
+
+        Parameters
+        ----------
+        seq_length : int, default=50
+            Sequence length for dummy input.
+        vocab_size : int, optional
+            Vocabulary size for generating dummy tokens.
+            If None, will try to get from model.vocab_size or model.item_num.
+        batch_size : int, default=2
+            Batch size for dummy input.
+        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.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        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.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+        ValueError
+            If vocab_size is not provided and cannot be inferred from model.
+
+        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
+
+        Examples
+        --------
+        >>> trainer = SeqTrainer(hstu_model, ...)
+        >>> trainer.fit(train_dl, val_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4, vocab_size=10000)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        try:
+            from torchview import draw_graph
+            TORCHVIEW_AVAILABLE = True
+        except ImportError:
+            TORCHVIEW_AVAILABLE = False
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "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"
+            )
+
+        from ..utils.visualization import _is_jupyter_environment, display_graph
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        # Get vocab_size from model if not provided
+        if vocab_size is None:
+            if hasattr(model, 'vocab_size'):
+                vocab_size = model.vocab_size
+            elif hasattr(model, 'item_num'):
+                vocab_size = model.item_num
+            else:
+                raise ValueError("vocab_size must be provided or model must have "
+                                 "'vocab_size' or 'item_num' attribute")
+
+        # Generate dummy inputs for sequence model
+        dummy_seq_tokens = torch.randint(0, vocab_size, (batch_size, seq_length), device=viz_device)
+        dummy_seq_time_diffs = torch.zeros(batch_size, seq_length, dtype=torch.float32, device=viz_device)
+
+        # Move model to device
+        model = model.to(viz_device)
+        model.eval()
+
+        # Call torchview.draw_graph
+        graph = draw_graph(model, input_data=(dummy_seq_tokens, dummy_seq_time_diffs), graph_name=graph_name, depth=depth, device=viz_device, expand_nested=expand_nested, show_shapes=show_shapes, save_graph=False, **kwargs)
+
+        # Set DPI for high-quality output
+        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('.')
+            output_format = ext if ext else 'pdf'
+            if directory != "." and not os.path.exists(directory):
+                os.makedirs(directory, exist_ok=True)
+            graph.visual_graph.render(filename=filename, directory=directory, format=output_format, cleanup=True)
+
+        # Handle default display behavior when save_path is None
+        if save_path is None:
+            if _is_jupyter_environment():
+                display_graph(graph)
+            else:
+                graph.visual_graph.view(cleanup=True)
+
+        return graph

torch_rechub/utils/model_utils.py

@@ -0,0 +1,233 @@
+"""Common model utility functions for Torch-RecHub.
+
+This module provides shared utilities for model introspection and input generation,
+used by both ONNX export and visualization features.
+
+Examples
+--------
+>>> from torch_rechub.utils.model_utils import extract_feature_info, generate_dummy_input
+>>> feature_info = extract_feature_info(model)
+>>> dummy_input = generate_dummy_input(feature_info['features'], batch_size=2)
+"""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+# Import feature types for type checking
+try:
+    from ..basic.features import DenseFeature, SequenceFeature, SparseFeature
+except ImportError:
+    # Fallback for standalone usage
+    SparseFeature = None
+    DenseFeature = None
+    SequenceFeature = None
+
+
+def extract_feature_info(model: nn.Module) -> Dict[str, Any]:
+    """Extract feature information from a torch-rechub model using reflection.
+
+    This function inspects model attributes to find feature lists without
+    modifying the model code. Supports various model architectures.
+
+    Parameters
+    ----------
+    model : nn.Module
+        The recommendation model to inspect.
+
+    Returns
+    -------
+    dict
+        Dictionary containing:
+        - 'features': List of unique Feature objects
+        - 'input_names': List of feature names in order
+        - 'input_types': Dict mapping feature name to feature type
+        - 'user_features': List of user-side features (for dual-tower models)
+        - 'item_features': List of item-side features (for dual-tower models)
+
+    Examples
+    --------
+    >>> from torch_rechub.models.ranking import DeepFM
+    >>> model = DeepFM(deep_features, fm_features, mlp_params)
+    >>> info = extract_feature_info(model)
+    >>> print(info['input_names'])  # ['user_id', 'item_id', ...]
+    """
+    # Common feature attribute names across different model types
+    feature_attrs = [
+        'features',           # MMOE, DCN, etc.
+        'deep_features',      # DeepFM, WideDeep
+        'fm_features',        # DeepFM
+        'wide_features',      # WideDeep
+        'linear_features',    # DeepFFM
+        'cross_features',     # DeepFFM
+        'user_features',      # DSSM, YoutubeDNN, MIND
+        'item_features',      # DSSM, YoutubeDNN, MIND
+        'history_features',   # DIN, MIND
+        'target_features',    # DIN
+        'neg_item_feature',   # YoutubeDNN, MIND
+    ]
+
+    all_features = []
+    user_features = []
+    item_features = []
+
+    for attr in feature_attrs:
+        if hasattr(model, attr):
+            feat_list = getattr(model, attr)
+            if isinstance(feat_list, list) and len(feat_list) > 0:
+                all_features.extend(feat_list)
+                # Track user/item features for dual-tower models
+                if 'user' in attr or 'history' in attr:
+                    user_features.extend(feat_list)
+                elif 'item' in attr:
+                    item_features.extend(feat_list)
+
+    # Deduplicate features by name while preserving order
+    seen = set()
+    unique_features = []
+    for f in all_features:
+        if hasattr(f, 'name') and f.name not in seen:
+            seen.add(f.name)
+            unique_features.append(f)
+
+    # Deduplicate user/item features
+    seen_user = set()
+    unique_user = [f for f in user_features if hasattr(f, 'name') and f.name not in seen_user and not seen_user.add(f.name)]
+    seen_item = set()
+    unique_item = [f for f in item_features if hasattr(f, 'name') and f.name not in seen_item and not seen_item.add(f.name)]
+
+    # Build input names and types
+    input_names = [f.name for f in unique_features if hasattr(f, 'name')]
+    input_types = {f.name: type(f).__name__ for f in unique_features if hasattr(f, 'name')}
+
+    return {
+        'features': unique_features,
+        'input_names': input_names,
+        'input_types': input_types,
+        'user_features': unique_user,
+        'item_features': unique_item,
+    }
+
+
+def generate_dummy_input(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Tuple[torch.Tensor, ...]:
+    """Generate dummy input tensors based on feature definitions.
+
+    Parameters
+    ----------
+    features : list
+        List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
+    batch_size : int, default=2
+        Batch size for dummy input.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature.
+    device : str, default='cpu'
+        Device to create tensors on.
+
+    Returns
+    -------
+    tuple of Tensor
+        Tuple of tensors in the order of input features.
+
+    Examples
+    --------
+    >>> features = [SparseFeature("user_id", 1000), SequenceFeature("hist", 500)]
+    >>> dummy = generate_dummy_input(features, batch_size=4)
+    >>> # Returns (user_id_tensor[4], hist_tensor[4, 10])
+    """
+    # Dynamic import to handle feature types
+    from ..basic.features import DenseFeature, SequenceFeature, SparseFeature
+
+    inputs = []
+    for feat in features:
+        if isinstance(feat, SequenceFeature):
+            # Sequence features have shape [batch_size, seq_length]
+            tensor = torch.randint(0, feat.vocab_size, (batch_size, seq_length), device=device)
+        elif isinstance(feat, SparseFeature):
+            # Sparse features have shape [batch_size]
+            tensor = torch.randint(0, feat.vocab_size, (batch_size,), device=device)
+        elif isinstance(feat, DenseFeature):
+            # Dense features always have shape [batch_size, embed_dim]
+            tensor = torch.randn(batch_size, feat.embed_dim, device=device)
+        else:
+            raise TypeError(f"Unsupported feature type: {type(feat)}")
+        inputs.append(tensor)
+    return tuple(inputs)
+
+
+def generate_dummy_input_dict(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Dict[str, torch.Tensor]:
+    """Generate dummy input dict based on feature definitions.
+
+    Similar to generate_dummy_input but returns a dict mapping feature names
+    to tensors. This is the expected input format for torch-rechub models.
+
+    Parameters
+    ----------
+    features : list
+        List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
+    batch_size : int, default=2
+        Batch size for dummy input.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature.
+    device : str, default='cpu'
+        Device to create tensors on.
+
+    Returns
+    -------
+    dict
+        Dict mapping feature names to tensors.
+
+    Examples
+    --------
+    >>> features = [SparseFeature("user_id", 1000)]
+    >>> dummy = generate_dummy_input_dict(features, batch_size=4)
+    >>> # Returns {"user_id": tensor[4]}
+    """
+    dummy_tuple = generate_dummy_input(features, batch_size, seq_length, device)
+    input_names = [f.name for f in features if hasattr(f, 'name')]
+    return {name: tensor for name, tensor in zip(input_names, dummy_tuple)}
+
+
+def generate_dynamic_axes(input_names: List[str], output_names: Optional[List[str]] = None, batch_dim: int = 0, include_seq_dim: bool = True, seq_features: Optional[List[str]] = None) -> Dict[str, Dict[int, str]]:
+    """Generate dynamic axes configuration for ONNX export.
+
+    Parameters
+    ----------
+    input_names : list of str
+        List of input tensor names.
+    output_names : list of str, optional
+        List of output tensor names. Default is ["output"].
+    batch_dim : int, default=0
+        Dimension index for batch size.
+    include_seq_dim : bool, default=True
+        Whether to include sequence dimension as dynamic.
+    seq_features : list of str, optional
+        List of feature names that are sequences.
+
+    Returns
+    -------
+    dict
+        Dynamic axes dict for torch.onnx.export.
+
+    Examples
+    --------
+    >>> axes = generate_dynamic_axes(["user_id", "item_id"], seq_features=["hist"])
+    >>> # Returns {"user_id": {0: "batch_size"}, "item_id": {0: "batch_size"}, ...}
+    """
+    if output_names is None:
+        output_names = ["output"]
+
+    dynamic_axes = {}
+
+    # Input axes
+    for name in input_names:
+        dynamic_axes[name] = {batch_dim: "batch_size"}
+        # Add sequence dimension for sequence features
+        if include_seq_dim and seq_features and name in seq_features:
+            dynamic_axes[name][1] = "seq_length"
+
+    # Output axes
+    for name in output_names:
+        dynamic_axes[name] = {batch_dim: "batch_size"}
+
+    return dynamic_axes

torch_rechub/utils/onnx_export.py

@@ -62,142 +62,9 @@ class ONNXWrapper(nn.Module):
             self.model.mode = self._original_mode
 
 
-def extract_feature_info(model: nn.Module) -> Dict[str, Any]:
-    """Extract feature information from a model using reflection.
-
-    This function inspects model attributes to find feature lists without
-    modifying the model code. Supports various model architectures.
-
-    Args:
-        model: The recommendation model to inspect.
-
-    Returns:
-        Dict containing:
-            - 'features': List of unique Feature objects
-            - 'input_names': List of feature names in order
-            - 'input_types': Dict mapping feature name to feature type
-            - 'user_features': List of user-side features (for dual-tower models)
-            - 'item_features': List of item-side features (for dual-tower models)
-    """
-    # Common feature attribute names across different model types
-    feature_attrs = [
-        'features',           # MMOE, DCN, etc.
-        'deep_features',      # DeepFM, WideDeep
-        'fm_features',        # DeepFM
-        'wide_features',      # WideDeep
-        'linear_features',    # DeepFFM
-        'cross_features',     # DeepFFM
-        'user_features',      # DSSM, YoutubeDNN, MIND
-        'item_features',      # DSSM, YoutubeDNN, MIND
-        'history_features',   # DIN, MIND
-        'target_features',    # DIN
-        'neg_item_feature',   # YoutubeDNN, MIND
-    ]
-
-    all_features = []
-    user_features = []
-    item_features = []
-
-    for attr in feature_attrs:
-        if hasattr(model, attr):
-            feat_list = getattr(model, attr)
-            if isinstance(feat_list, list) and len(feat_list) > 0:
-                all_features.extend(feat_list)
-                # Track user/item features for dual-tower models
-                if 'user' in attr or 'history' in attr:
-                    user_features.extend(feat_list)
-                elif 'item' in attr:
-                    item_features.extend(feat_list)
-
-    # Deduplicate features by name while preserving order
-    seen = set()
-    unique_features = []
-    for f in all_features:
-        if hasattr(f, 'name') and f.name not in seen:
-            seen.add(f.name)
-            unique_features.append(f)
-
-    # Deduplicate user/item features
-    seen_user = set()
-    unique_user = [f for f in user_features if hasattr(f, 'name') and f.name not in seen_user and not seen_user.add(f.name)]
-    seen_item = set()
-    unique_item = [f for f in item_features if hasattr(f, 'name') and f.name not in seen_item and not seen_item.add(f.name)]
-
-    return {
-        'features': unique_features,
-        'input_names': [f.name for f in unique_features],
-        'input_types': {
-            f.name: type(f).__name__ for f in unique_features
-        },
-        'user_features': unique_user,
-        'item_features': unique_item,
-    }
-
-
-def generate_dummy_input(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Tuple[torch.Tensor, ...]:
-    """Generate dummy input tensors for ONNX export based on feature definitions.
-
-    Args:
-        features: List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
-        batch_size: Batch size for dummy input (default: 2).
-        seq_length: Sequence length for SequenceFeature (default: 10).
-        device: Device to create tensors on (default: 'cpu').
-
-    Returns:
-        Tuple of tensors in the order of input features.
-
-    Example:
-        >>> features = [SparseFeature("user_id", 1000), SequenceFeature("hist", 500)]
-        >>> dummy = generate_dummy_input(features, batch_size=4)
-        >>> # Returns (user_id_tensor[4], hist_tensor[4, 10])
-    """
-    inputs = []
-    for feat in features:
-        if isinstance(feat, SequenceFeature):
-            # Sequence features have shape [batch_size, seq_length]
-            tensor = torch.randint(0, feat.vocab_size, (batch_size, seq_length), device=device)
-        elif isinstance(feat, SparseFeature):
-            # Sparse features have shape [batch_size]
-            tensor = torch.randint(0, feat.vocab_size, (batch_size,), device=device)
-        elif isinstance(feat, DenseFeature):
-            # Dense features have shape [batch_size, embed_dim]
-            tensor = torch.randn(batch_size, feat.embed_dim, device=device)
-        else:
-            raise TypeError(f"Unsupported feature type: {type(feat)}")
-        inputs.append(tensor)
-    return tuple(inputs)
-
-
-def generate_dynamic_axes(input_names: List[str], output_names: List[str] = None, batch_dim: int = 0, include_seq_dim: bool = True, seq_features: List[str] = None) -> Dict[str, Dict[int, str]]:
-    """Generate dynamic axes configuration for ONNX export.
-
-    Args:
-        input_names: List of input tensor names.
-        output_names: List of output tensor names (default: ["output"]).
-        batch_dim: Dimension index for batch size (default: 0).
-        include_seq_dim: Whether to include sequence dimension as dynamic (default: True).
-        seq_features: List of feature names that are sequences (default: auto-detect).
-
-    Returns:
-        Dynamic axes dict for torch.onnx.export.
-    """
-    if output_names is None:
-        output_names = ["output"]
-
-    dynamic_axes = {}
-
-    # Input axes
-    for name in input_names:
-        dynamic_axes[name] = {batch_dim: "batch_size"}
-        # Add sequence dimension for sequence features
-        if include_seq_dim and seq_features and name in seq_features:
-            dynamic_axes[name][1] = "seq_length"
-
-    # Output axes
-    for name in output_names:
-        dynamic_axes[name] = {batch_dim: "batch_size"}
-
-    return dynamic_axes
+# 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:

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

{torch_rechub-0.0.4.dist-info → torch_rechub-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-rechub
-Version: 0.0.4
+Version: 0.0.5
 Summary: A Pytorch Toolbox for Recommendation Models, Easy-to-use and Easy-to-extend.
 Project-URL: Homepage, https://github.com/datawhalechina/torch-rechub
 Project-URL: Documentation, https://www.torch-rechub.com
@@ -41,6 +41,9 @@ Requires-Dist: yapf==0.43.0; extra == 'dev'
 Provides-Extra: onnx
 Requires-Dist: onnx>=1.12.0; extra == 'onnx'
 Requires-Dist: onnxruntime>=1.12.0; extra == 'onnx'
+Provides-Extra: visualization
+Requires-Dist: graphviz>=0.20; extra == 'visualization'
+Requires-Dist: torchview>=0.2.6; extra == 'visualization'
 Description-Content-Type: text/markdown
 
 # 🔥 Torch-RecHub - 轻量、高效、易用的 PyTorch 推荐系统框架
@@ -69,13 +72,13 @@ Description-Content-Type: text/markdown
 
 ## 🎯 为什么选择 Torch-RecHub？
 
-| 特性 | Torch-RecHub | 其他框架 |
-|------|-------------|---------|
-| 代码行数 | **10行** 完成训练+评估+部署 | 100+ 行 |
-| 模型覆盖 | **30+** 主流模型 | 有限 |
-| 生成式推荐 | ✅ HSTU/HLLM (Meta 2024) | ❌ |
-| ONNX 一键导出 | ✅ 内置支持 | 需手动适配 |
-| 学习曲线 | 极低 | 陡峭 |
+| 特性          | Torch-RecHub                | 其他框架   |
+| ------------- | --------------------------- | ---------- |
+| 代码行数      | **10行** 完成训练+评估+部署 | 100+ 行    |
+| 模型覆盖      | **30+** 主流模型            | 有限       |
+| 生成式推荐    | ✅ HSTU/HLLM (Meta 2024)     | ❌          |
+| ONNX 一键导出 | ✅ 内置支持                  | 需手动适配 |
+| 学习曲线      | 极低                        | 陡峭       |
 
 ## ✨ 特性
 
@@ -205,52 +208,52 @@ torch-rechub/             # 根目录
 
 ### 排序模型 (Ranking Models) - 13个
 
-| 模型 | 论文 | 简介 |
-|------|------|------|
-| **DeepFM** | [IJCAI 2017](https://arxiv.org/abs/1703.04247) | FM + Deep 联合训练 |
-| **Wide&Deep** | [DLRS 2016](https://arxiv.org/abs/1606.07792) | 记忆 + 泛化能力结合 |
-| **DCN** | [KDD 2017](https://arxiv.org/abs/1708.05123) | 显式特征交叉网络 |
-| **DCN-v2** | [WWW 2021](https://arxiv.org/abs/2008.13535) | 增强版交叉网络 |
-| **DIN** | [KDD 2018](https://arxiv.org/abs/1706.06978) | 注意力机制捕捉用户兴趣 |
-| **DIEN** | [AAAI 2019](https://arxiv.org/abs/1809.03672) | 兴趣演化建模 |
-| **BST** | [DLP-KDD 2019](https://arxiv.org/abs/1905.06874) | Transformer 序列建模 |
-| **AFM** | [IJCAI 2017](https://arxiv.org/abs/1708.04617) | 注意力因子分解机 |
-| **AutoInt** | [CIKM 2019](https://arxiv.org/abs/1810.11921) | 自动特征交互学习 |
-| **FiBiNET** | [RecSys 2019](https://arxiv.org/abs/1905.09433) | 特征重要性 + 双线性交互 |
-| **DeepFFM** | [RecSys 2019](https://arxiv.org/abs/1611.00144) | 场感知因子分解机 |
-| **EDCN** | [KDD 2021](https://arxiv.org/abs/2106.03032) | 增强型交叉网络 |
+| 模型          | 论文                                             | 简介                    |
+| ------------- | ------------------------------------------------ | ----------------------- |
+| **DeepFM**    | [IJCAI 2017](https://arxiv.org/abs/1703.04247)   | FM + Deep 联合训练      |
+| **Wide&Deep** | [DLRS 2016](https://arxiv.org/abs/1606.07792)    | 记忆 + 泛化能力结合     |
+| **DCN**       | [KDD 2017](https://arxiv.org/abs/1708.05123)     | 显式特征交叉网络        |
+| **DCN-v2**    | [WWW 2021](https://arxiv.org/abs/2008.13535)     | 增强版交叉网络          |
+| **DIN**       | [KDD 2018](https://arxiv.org/abs/1706.06978)     | 注意力机制捕捉用户兴趣  |
+| **DIEN**      | [AAAI 2019](https://arxiv.org/abs/1809.03672)    | 兴趣演化建模            |
+| **BST**       | [DLP-KDD 2019](https://arxiv.org/abs/1905.06874) | Transformer 序列建模    |
+| **AFM**       | [IJCAI 2017](https://arxiv.org/abs/1708.04617)   | 注意力因子分解机        |
+| **AutoInt**   | [CIKM 2019](https://arxiv.org/abs/1810.11921)    | 自动特征交互学习        |
+| **FiBiNET**   | [RecSys 2019](https://arxiv.org/abs/1905.09433)  | 特征重要性 + 双线性交互 |
+| **DeepFFM**   | [RecSys 2019](https://arxiv.org/abs/1611.00144)  | 场感知因子分解机        |
+| **EDCN**      | [KDD 2021](https://arxiv.org/abs/2106.03032)     | 增强型交叉网络          |
 
 ### 召回模型 (Matching Models) - 12个
 
-| 模型 | 论文 | 简介 |
-|------|------|------|
-| **DSSM** | [CIKM 2013](https://posenhuang.github.io/papers/cikm2013_DSSM_fullversion.pdf) | 经典双塔召回模型 |
-| **YoutubeDNN** | [RecSys 2016](https://dl.acm.org/doi/10.1145/2959100.2959190) | YouTube 深度召回 |
-| **YoutubeSBC** | [RecSys 2019](https://dl.acm.org/doi/10.1145/3298689.3346997) | 采样偏差校正版本 |
-| **MIND** | [CIKM 2019](https://arxiv.org/abs/1904.08030) | 多兴趣动态路由 |
-| **SINE** | [WSDM 2021](https://arxiv.org/abs/2103.06920) | 稀疏兴趣网络 |
-| **GRU4Rec** | [ICLR 2016](https://arxiv.org/abs/1511.06939) | GRU 序列推荐 |
-| **SASRec** | [ICDM 2018](https://arxiv.org/abs/1808.09781) | 自注意力序列推荐 |
-| **NARM** | [CIKM 2017](https://arxiv.org/abs/1711.04725) | 神经注意力会话推荐 |
-| **STAMP** | [KDD 2018](https://dl.acm.org/doi/10.1145/3219819.3219895) | 短期注意力记忆优先 |
-| **ComiRec** | [KDD 2020](https://arxiv.org/abs/2005.09347) | 可控多兴趣推荐 |
+| 模型           | 论文                                                                           | 简介               |
+| -------------- | ------------------------------------------------------------------------------ | ------------------ |
+| **DSSM**       | [CIKM 2013](https://posenhuang.github.io/papers/cikm2013_DSSM_fullversion.pdf) | 经典双塔召回模型   |
+| **YoutubeDNN** | [RecSys 2016](https://dl.acm.org/doi/10.1145/2959100.2959190)                  | YouTube 深度召回   |
+| **YoutubeSBC** | [RecSys 2019](https://dl.acm.org/doi/10.1145/3298689.3346997)                  | 采样偏差校正版本   |
+| **MIND**       | [CIKM 2019](https://arxiv.org/abs/1904.08030)                                  | 多兴趣动态路由     |
+| **SINE**       | [WSDM 2021](https://arxiv.org/abs/2103.06920)                                  | 稀疏兴趣网络       |
+| **GRU4Rec**    | [ICLR 2016](https://arxiv.org/abs/1511.06939)                                  | GRU 序列推荐       |
+| **SASRec**     | [ICDM 2018](https://arxiv.org/abs/1808.09781)                                  | 自注意力序列推荐   |
+| **NARM**       | [CIKM 2017](https://arxiv.org/abs/1711.04725)                                  | 神经注意力会话推荐 |
+| **STAMP**      | [KDD 2018](https://dl.acm.org/doi/10.1145/3219819.3219895)                     | 短期注意力记忆优先 |
+| **ComiRec**    | [KDD 2020](https://arxiv.org/abs/2005.09347)                                   | 可控多兴趣推荐     |
 
 ### 多任务模型 (Multi-Task Models) - 5个
 
-| 模型 | 论文 | 简介 |
-|------|------|------|
-| **ESMM** | [SIGIR 2018](https://arxiv.org/abs/1804.07931) | 全空间多任务建模 |
-| **MMoE** | [KDD 2018](https://dl.acm.org/doi/10.1145/3219819.3220007) | 多门控专家混合 |
-| **PLE** | [RecSys 2020](https://dl.acm.org/doi/10.1145/3383313.3412236) | 渐进式分层提取 |
-| **AITM** | [KDD 2021](https://arxiv.org/abs/2105.08489) | 自适应信息迁移 |
-| **SharedBottom** | - | 经典多任务共享底层 |
+| 模型             | 论文                                                          | 简介               |
+| ---------------- | ------------------------------------------------------------- | ------------------ |
+| **ESMM**         | [SIGIR 2018](https://arxiv.org/abs/1804.07931)                | 全空间多任务建模   |
+| **MMoE**         | [KDD 2018](https://dl.acm.org/doi/10.1145/3219819.3220007)    | 多门控专家混合     |
+| **PLE**          | [RecSys 2020](https://dl.acm.org/doi/10.1145/3383313.3412236) | 渐进式分层提取     |
+| **AITM**         | [KDD 2021](https://arxiv.org/abs/2105.08489)                  | 自适应信息迁移     |
+| **SharedBottom** | -                                                             | 经典多任务共享底层 |
 
 ### 生成式推荐 (Generative Recommendation) - 2个
 
-| 模型 | 论文 | 简介 |
-|------|------|------|
+| 模型     | 论文                                          | 简介                                         |
+| -------- | --------------------------------------------- | -------------------------------------------- |
 | **HSTU** | [Meta 2024](https://arxiv.org/abs/2402.17152) | 层级序列转换单元，支撑 Meta 万亿参数推荐系统 |
-| **HLLM** | [2024](https://arxiv.org/abs/2409.12740) | 层级大语言模型推荐，融合 LLM 语义理解能力 |
+| **HLLM** | [2024](https://arxiv.org/abs/2409.12740)      | 层级大语言模型推荐，融合 LLM 语义理解能力    |
 
 ## 📊 支持的数据集
 
@@ -338,11 +341,19 @@ model = DSSM(user_features, item_features, temperature=0.02,
 match_trainer = MatchTrainer(model)
 match_trainer.fit(train_dl)
 match_trainer.export_onnx("dssm.onnx")
-# 双塔模型可分别导出用户塔和物品塔: 
+# 双塔模型可分别导出用户塔和物品塔:
 # match_trainer.export_onnx("user_tower.onnx", mode="user")
 # match_trainer.export_onnx("dssm_item.onnx", tower="item")
 ```
 
+### 模型可视化
+
+```python
+# 可视化模型架构（需要安装: pip install torch-rechub[visualization]）
+graph = ctr_trainer.visualization(depth=4)  # 生成计算图
+ctr_trainer.visualization(save_path="model.pdf", dpi=300)  # 保存为高清 PDF
+```
+
 ## 👨‍💻‍ 贡献者
 
 感谢所有的贡献者！

{torch_rechub-0.0.4.dist-info → torch_rechub-0.0.5.dist-info}/RECORD

@@ -45,18 +45,20 @@ torch_rechub/models/ranking/edcn.py,sha256=6f_S8I6Ir16kCIU54R4EfumWfUFOND5KDKUPH
 torch_rechub/models/ranking/fibinet.py,sha256=fmEJ9WkO8Mn0RtK_8aRHlnQFh_jMBPO0zODoHZPWmDA,2234
 torch_rechub/models/ranking/widedeep.py,sha256=eciRvWRBHLlctabLLS5NB7k3MnqrWXCBdpflOU6jMB0,1636
 torch_rechub/trainers/__init__.py,sha256=NSa2DqgfE1HGDyj40YgrbtUrfBHBxNBpw57XtaAB_jE,148
-torch_rechub/trainers/ctr_trainer.py,sha256=RDUXkn7GwLzs3f0kWZwGDNCpqiMeGXo7R6ezFeZdPg8,9075
-torch_rechub/trainers/match_trainer.py,sha256=xox5eaPKjSgErJQpbSr29sbyGs1p2sFaKEjxACE6uMI,11276
+torch_rechub/trainers/ctr_trainer.py,sha256=ECXaK0x2_6jZVxtEazgN3hkBpSAMPeGeNtunqI_OECo,12860
+torch_rechub/trainers/match_trainer.py,sha256=QHZb32Rf7yp-NvEzdeiG1HQghQ76_vuu59K1IsdK60k,15055
 torch_rechub/trainers/matching.md,sha256=vIBQ3UMmVpUpyk38rrkelFwm_wXVXqMOuqzYZ4M8bzw,30
-torch_rechub/trainers/mtl_trainer.py,sha256=tC4c2KIc-H8Wvj4qCzcW6TyfMLRPJyfQvTaN0dDePFg,12598
-torch_rechub/trainers/seq_trainer.py,sha256=lXKRx7XbZ3iJuqp_f05vw_jkn8X5j8HmH6Nr-typiIU,12043
+torch_rechub/trainers/mtl_trainer.py,sha256=MjasE_QOPfGxiUW1JpYYQ2iuBSSk-lissAGp4Sw1CWk,16427
+torch_rechub/trainers/seq_trainer.py,sha256=uAo9XymwQupCqvm5otKW81tz1nxd3crJ2ul2r7lrEAE,17633
 torch_rechub/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 torch_rechub/utils/data.py,sha256=vzLAAVt6dujg_vbGhQewiJc0l6JzwzdcM_9EjoOz898,19882
 torch_rechub/utils/hstu_utils.py,sha256=qLON_pJDC-kDyQn1PoN_HaHi5xTNCwZPgJeV51Z61Lc,6207
 torch_rechub/utils/match.py,sha256=l9qDwJGHPP9gOQTMYoqGVdWrlhDx1F1-8UnQwDWrEyk,18143
+torch_rechub/utils/model_utils.py,sha256=VLhSbTpupxrFyyY3NzMQ32PPmo5YHm1T96u9KDlwiWE,8450
 torch_rechub/utils/mtl.py,sha256=AxU05ezizCuLdbPuCg1ZXE0WAStzuxaS5Sc3nwMCBpI,5737
-torch_rechub/utils/onnx_export.py,sha256=uRcAD4uZ3eIQbM-DPhdc0bkaPaslNsOYny6BOeLVBfU,13660
-torch_rechub-0.0.4.dist-info/METADATA,sha256=SNm71v_YOfculnc13p266bD_8yLo0U_16F_aJQPDvYo,16149
-torch_rechub-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-torch_rechub-0.0.4.dist-info/licenses/LICENSE,sha256=V7ietiX9G_84HtgEbxDgxClniqXGm2t5q8WM4AHGTu0,1066
-torch_rechub-0.0.4.dist-info/RECORD,,
+torch_rechub/utils/onnx_export.py,sha256=LRHyZaR9zZJyg6xtuqQHWmusWq-yEvw9EhlmoEwcqsg,8364
+torch_rechub/utils/visualization.py,sha256=Djv8W5SkCk3P2dol5VXf0_eanIhxDwRd7fzNOQY4uiU,9506
+torch_rechub-0.0.5.dist-info/METADATA,sha256=7k9N1xGB4JeWzri7iA7kJbPnAJ-KhXF7vBV-_b8Ghrg,17998
+torch_rechub-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+torch_rechub-0.0.5.dist-info/licenses/LICENSE,sha256=V7ietiX9G_84HtgEbxDgxClniqXGm2t5q8WM4AHGTu0,1066
+torch_rechub-0.0.5.dist-info/RECORD,,