model-sketch 0.0.4

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