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:

Customization Options

Customize model visualizations with the following SketchOptions:

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.