Visualization

Torch-RecHub provides model architecture visualization to help developers intuitively understand model structure and computation flow.

Why torchview

Torch-RecHub uses torchview as its visualization backend. The project's wrapper passes the feature dictionary to the model as one positional argument, so it can trace the dictionary inputs commonly used by recommendation models.

Recommendation model inputs are typically dictionaries containing multiple feature types:

# Typical input format for recommendation models
x = {
    "user_id": tensor([1, 2, 3]),           # Sparse feature
    "age": tensor([0.5, 0.3, 0.8]),         # Dense feature
    "hist_items": tensor([[1,2,3], ...]),   # Sequence feature
}
model(x)  # Dictionary as input

torchviz is based on the autograd graph, while Netron primarily inspects an already exported static model. These tools serve different purposes; the distinction cannot be reduced to whether they "support dictionaries."

Tip: If you have exported your model to ONNX format, you can also use Netron to view the model structure online. See ONNX Export Documentation.

Feature	torchview	torchviz	netron
This project's dictionary-input wrapper	✅	Requires a custom wrapper	Requires model export first
Forward pass tracing	✅	❌ (autograd-based)	❌ (static parsing)
Shows one actually executed path	✅	✅	❌ (static parsing)
Show tensor shapes	✅	❌	✅
Adjustable depth	✅	❌	❌
Nested module expansion	✅	❌	Partial

Other advantages:

• Forward tracing: Captures the modules traversed by this input during one forward pass; control-flow branches that are not executed do not appear in the graph
• Depth control: Flexibly control display granularity via depth parameter
• Shape visualization: Intuitively display tensor shapes at each layer

Installation

Visualization requires additional dependencies:

pip install "torch-rechub[visualization]"

Also install system-level graphviz:

# Ubuntu/Debian
sudo apt-get install graphviz

# macOS
brew install graphviz

# Windows
choco install graphviz

Quick Start

from torch_rechub.models.ranking import DeepFM
from torch_rechub.utils.visualization import visualize_model

# Create model
model = DeepFM(
    deep_features=deep_features,
    fm_features=fm_features,
    mlp_params={"dims": [256, 128], "dropout": 0.2}
)

# Visualize (auto-display in Jupyter)
graph = visualize_model(model, depth=4)

# Save as PDF
visualize_model(model, save_path="model_arch.pdf", dpi=300)

Core Functions

visualize_model

Generate computation graph visualization.

from torch_rechub.utils.visualization import visualize_model

graph = visualize_model(
    model,                    # PyTorch model
    input_data=None,          # Input data (optional, auto-generated)
    batch_size=2,             # Batch size for auto-generated input
    seq_length=10,            # Sequence feature length
    depth=3,                  # Visualization depth
    show_shapes=True,         # Show tensor shapes
    expand_nested=True,       # Expand nested modules
    save_path=None,           # Save path
    graph_name="model",       # Graph name
    device="cpu",             # Device
    dpi=300,                  # Output resolution
)

Parameter Description:

Parameter	Type	Description	Default
model	nn.Module	PyTorch model to visualize	Required
input_data	dict	Input data dictionary, auto-generated if None	None
batch_size	int	Batch size for auto-generated input	2
seq_length	int	Sequence feature length	10
depth	int	Visualization depth, -1 shows all layers	3
show_shapes	bool	Show tensor shapes on edges	True
expand_nested	bool	Expand nested nn.Module	True
save_path	str	Save path, supports .pdf/.svg/.png	None
dpi	int	Output image resolution	300

display_graph

Display computation graph in Jupyter.

from torch_rechub.utils.visualization import display_graph

# Get the computation graph
graph = visualize_model(model, depth=4)

# Display it in Jupyter
display_graph(graph, format='png')

Usage Examples

Ranking Model Visualization

from torch_rechub.models.ranking import DeepFM
from torch_rechub.utils.visualization import visualize_model
from torch_rechub.basic.features import DenseFeature, SparseFeature

# Define features
dense_features = [DenseFeature("age"), DenseFeature("income")]
sparse_features = [
    SparseFeature("city", vocab_size=100, embed_dim=16),
    # Sparse features used in FM interactions must share the same embedding dimension
    SparseFeature("gender", vocab_size=3, embed_dim=16)
]

# Create model
model = DeepFM(
    deep_features=sparse_features + dense_features,
    fm_features=sparse_features,
    mlp_params={"dims": [256, 128, 64], "dropout": 0.2}
)

# Visualize
visualize_model(model, depth=4, save_path="deepfm_arch.pdf")

Retrieval Model Visualization

from torch_rechub.models.matching import DSSM
from torch_rechub.utils.visualization import visualize_model

model = DSSM(
    user_features=user_features,
    item_features=item_features,
    temperature=0.02,
    user_params={"dims": [256, 128, 64]},
    item_params={"dims": [256, 128, 64]}
)

visualize_model(model, depth=3, save_path="dssm_arch.png", dpi=300)

Visualization via Trainer

Trainers also provide a visualization method:

from torch_rechub.trainers import CTRTrainer

trainer = CTRTrainer(model)
trainer.fit(train_dl, val_dl)

# Visualize model
trainer.visualization(save_path="model.pdf", depth=4)

Output Formats

Format	Extension	Use Case
PDF	.pdf	Papers, reports (vector, scalable)
SVG	.svg	Web display (vector)
PNG	.png	General image format

FAQ

Q: graphviz not installed error?

Make sure both Python package and system package are installed:

pip install graphviz
# Plus system-level installation (see above)

Q: Image not displaying in VSCode?

Try setting output format to PNG:

import graphviz
graphviz.set_jupyter_format('png')

Q: How to adjust image size?

Use the dpi parameter to control resolution, or use the returned graph object:

graph = visualize_model(model, depth=4)
graph.resize_graph(scale=1.5)