model-sketch 0.0.12

Publication-quality visualization of PyTorch model architectures

model-sketch

Image generated by AI

Python package for generating publication-quality visualizations of PyTorch model architectures. It focuses on clean SVG output, optional PNG export, and a simple API for sketching models from example forward inputs.

Installation

Base installation:

pip install model-sketch

If you also want PNG export support (CairoSVG backend), install the optional PNG extras:

pip install model-sketch[png]

Examples

Explore practical examples demonstrating model-sketch features:

Example	Description	Notebook	Colab
Tiny MLP sketch	Build a small MLP, render SVG/PNG, and save outputs to files.	simple_model.ipynb	Open in Colab
Semantic options demo	Compare semantic_view='blocks' vs semantic_view='layers' and semantic spacing on a larger model.	semantic_options.ipynb	Open in Colab
Edge style options demo	Compare edge_style='straight' and edge_style='orthogonal' with arrowheads on a branching model.	edge_style_options.ipynb	Open in Colab
Repeated patterns demo	Inspect repeated motif behavior across layers and blocks views, including repeated highlights and collapsed motifs.	repeated_patterns_demo.ipynb	Open in Colab
Model diff demo	Compare baseline and candidate models, then inspect added/removed/shared architecture labels.	model_diff_demo.ipynb	Open in Colab
Architecture coverage demo	Run CNN, BiLSTM, BiGRU, Transformer, multi-input, dict-output, and UNet-like examples through the same sketch pipeline.	architecture_coverage.ipynb	Open in Colab
Customization options demo	Demonstrate layer/block renaming and coloring, shape text toggling, and orientation choices.	customization_options_demo.ipynb	Open in Colab

Customization Options

Customize model visualizations with the following SketchOptions:

Option	Type	Default	Description
show_shape_text	bool	True	Show shape information (e.g., [2, 3]) on layer/block labels.
layer_name_map	dict	None	Rename layers via {"layer_id": "custom name"}. Supports both internal IDs and displayed labels.
block_name_map	dict	None	Rename blocks (when semantic_view='blocks') via {"block_id": "custom name"}. Supports "Inputs"/"Outputs" friendly labels.
layer_color_map	dict	None	Set layer fill colors via {"layer_id": "#hex"} in layers view.
block_color_map	dict	None	Set block fill colors via {"block_id": "#hex"} in blocks view.
semantic_grouping	str	"module_prefix"	Block grouping strategy: "module_prefix" or "single_block".
semantic_inter_block_gap	float	96.0	Extra spacing between semantic blocks in layout planning.
semantic_view	str	"layers"	Show layers ("layers") or semantic blocks ("blocks").
layout_orientation	str	"ltr"	Layout direction: "ltr" (left-to-right), "rtl", "ttb" (top-to-bottom), or "btt".
show_residual_highlight	bool	True	Highlight residual/skip edges when detected.
show_repeated_pattern_highlight	bool	False	Highlight repeated-pattern nodes when detected.
edge_style	str	"straight"	Edge rendering style: "straight" or "orthogonal".
edge_arrowheads	bool	False	Show arrowheads on edges to indicate data flow direction.

Quick Examples

Hide shape information:

from modelsketch import SketchOptions, sketch_model

result = sketch_model(
    model,
    example_args=(example_input,),
    options=SketchOptions(show_shape_text=False),
)
svg = result.svg

Rename layers:

result = sketch_model(
    model,
    example_args=(example_input,),
    options=SketchOptions(
        layer_name_map={"fc1": "Dense Layer 1", "fc2": "Output Layer"}
    ),
)
svg = result.svg

Rename input/output blocks and customize blocks view:

result = sketch_model(
    model,
    example_args=(example_input,),
    options=SketchOptions(
        semantic_view="blocks",
        block_name_map={"Inputs": "Model Inputs", "stage1": "Early Layers"},
    ),
)
svg = result.svg

Add arrowheads and use vertical layout:

result = sketch_model(
    model,
    example_args=(example_input,),
    options=SketchOptions(
        edge_arrowheads=True,
        layout_orientation="ttb",
    ),
)
svg = result.svg

Apply custom layer colors:

result = sketch_model(
    model,
    example_args=(example_input,),
    options=SketchOptions(
        semantic_view="layers",
        layer_color_map={"stem": "#e0f4ff", "head": "#ffe0e0"},
    ),
)
svg = result.svg

For comprehensive examples covering all options, see customization_options_demo.ipynb.