Framework-agnostic neural network visualization for Jupyter notebooks
Documentation • Features • Installation • Quick Start • Examples • 3D Visualization • API • Contributing
modelviz generates beautiful, publication-ready neural network architecture diagrams from your PyTorch and TensorFlow/Keras models. Simply pass your model object and get a stunning visualization — no manual diagram creation required.
| Feature | Description |
|---|---|
| 🔍 Auto-detection | Automatically detects PyTorch and TensorFlow/Keras models |
| 📊 2D Diagrams | Clean Graphviz diagrams with layer types, shapes, and parameters |
| 🎮 3D Interactive | Stunning Three.js visualizations with distinct shapes per layer |
| 🔄 Skip Connections | ResNet-style residual paths, dense connections, and branching architectures |
| 🎨 Smart Styling | Color-coded nodes for Conv, Linear, Pooling, Activation layers |
| 📦 Block Grouping | Auto-merges common patterns (Conv+ReLU, Conv+BN+ReLU) |
| 📓 Notebook-native | Renders inline in Jupyter, Colab, and VSCode notebooks |
| 💾 Export | Save as PNG, SVG, PDF, or interactive HTML |
demo.-.modelviz-ai.mov
Each layer type has a distinct, meaningful 3D representation:
| Layer | Shape | Rationale |
|---|---|---|
| Conv2d | 3D Box | Feature maps are 3D volumes (C×H×W) |
| Linear | Flat Plane | Weight matrix is 2D |
| Pooling | Small Cube | Reduces spatial dimensions |
| Activation | Sphere | Element-wise uniform operation |
| BatchNorm | Thin Slab | Normalizes distribution |
| Flatten | Cone | Funnels data to 1D |
| Dropout | Wireframe | Sparse/dropped neurons |
| RNN/LSTM | Cylinder | Recurrent/cyclical flow |
| Attention | Octahedron | Multi-head patterns |
# Basic installation
pip install modelviz-ai
# With PyTorch support
pip install modelviz-ai[torch]
# With TensorFlow support
pip install modelviz-ai[tf]
# All frameworks + development tools
pip install modelviz-ai[all,dev]git clone https://github.com/shreyanshjain05/modelviz.git
cd modelviz
pip install -e ".[dev]"For 2D Graphviz diagrams, install the Graphviz system package:
# macOS
brew install graphviz
# Ubuntu/Debian
sudo apt-get install graphviz
# Windows (or use Conda)
conda install -c conda-forge graphvizNote: Three.js 3D visualizations work without any system dependencies.
import torch.nn as nn
from modelviz import visualize
model = nn.Sequential(
nn.Conv2d(1, 32, 3),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(32 * 13 * 13, 10)
)
# Renders inline in Jupyter
visualize(model, input_shape=(1, 1, 28, 28))
# Save to file
visualize(model, input_shape=(1, 1, 28, 28), save_path="model.png")from modelviz import visualize_threejs
# Creates an interactive HTML file
visualize_threejs(
model,
input_shape=(1, 1, 28, 28),
save_path="model_3d.html"
)
# Open model_3d.html in your browser!import torch.nn as nn
from modelviz import visualize, visualize_threejs
class CNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 8 * 8, 256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, 10),
)
def forward(self, x):
return self.classifier(self.features(x))
model = CNN()
# 2D diagram with layer grouping
visualize(model, input_shape=(1, 3, 32, 32), title="CNN Architecture")
# 3D interactive visualization
visualize_threejs(model, input_shape=(1, 3, 32, 32), save_path="cnn_3d.html")import tensorflow as tf
from modelviz import visualize
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(28, 28, 1)),
tf.keras.layers.Conv2D(32, 3, activation='relu'),
tf.keras.layers.MaxPooling2D(2),
tf.keras.layers.Conv2D(64, 3, activation='relu'),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10, activation='softmax'),
])
# No input_shape needed - Keras models are already built
visualize(model, save_path="keras_model.svg")The Three.js renderer creates stunning interactive 3D diagrams:
from modelviz import visualize_threejs
html = visualize_threejs(
model,
input_shape=(1, 3, 224, 224),
title="ResNet Block",
show_shapes=True, # Show tensor dimensions
show_params=True, # Show parameter counts
group_blocks=True, # Merge Conv+BN+ReLU
save_path="resnet.html"
)| Action | Control |
|---|---|
| Rotate | Drag mouse |
| Zoom | Scroll wheel |
| Pan | Shift + Drag |
| Details | Hover over layer |
- Horizontal layout — Data flows left to right
- Text labels — Layer type and output shape above each node
- Animated particles — Shows data flow between layers
- Hover tooltips — Full layer information on mouseover
- Legend — Color and shape guide
Generate a 2D Graphviz diagram.
visualize(
model, # PyTorch or Keras model
input_shape=(1, 3, 224, 224), # Required for PyTorch
framework="auto", # "auto", "pytorch", "tensorflow"
show_shapes=True, # Show output tensor shapes
show_params=True, # Show parameter counts
group_blocks=True, # Merge Conv+ReLU patterns
save_path="model.png", # Optional: save to file
title="My Model", # Optional: diagram title
) -> graphviz.DigraphGenerate an interactive 3D Three.js visualization.
visualize_threejs(
model, # PyTorch or Keras model
input_shape=(1, 3, 224, 224), # Required for PyTorch
framework="auto", # "auto", "pytorch", "tensorflow"
show_shapes=True, # Show shapes in labels
show_params=True, # Show params in tooltips
group_blocks=True, # Merge Conv+ReLU patterns
save_path="model.html", # Save as HTML file
title="My Model 3D", # Visualization title
) -> str # Returns HTML stringGenerate a Plotly 3D visualization (simpler fallback).
visualize_3d(
model,
input_shape=(1, 3, 224, 224),
layout="tower", # "tower", "spiral", "grid"
save_path="model.png",
) -> plotly.graph_objects.Figure| Layer Type | Color | Hex |
|---|---|---|
| Convolution | Indigo | #6366f1 |
| Linear/Dense | Purple | #8b5cf6 |
| Pooling | Cyan | #06b6d4 |
| Activation | Amber | #f59e0b |
| Normalization | Emerald | #10b981 |
| Flatten | Pink | #ec4899 |
| Dropout | Red | #ef4444 |
| Embedding | Lime | #84cc16 |
| RNN/LSTM | Teal | #14b8a6 |
| Attention | Orange | #f97316 |
Common patterns are automatically merged:
Conv2d→BatchNorm2d→ReLU→ Conv2d + BatchNorm2d + ReLUConv2d→ReLU→ Conv2d + ReLULinear→ReLU→ Linear + ReLUDense→Activation→ Dense + Activation
Disable with group_blocks=False.
modelviz/
├── modelviz/
│ ├── __init__.py # Public API
│ ├── visualize.py # Main API functions
│ ├── graph/
│ │ ├── layer_node.py # LayerNode dataclass
│ │ └── builder.py # Graph construction
│ ├── parsers/
│ │ ├── torch_parser.py # PyTorch model parsing
│ │ ├── tf_parser.py # TensorFlow/Keras parsing
│ │ └── fx_tracer.py # Skip connection detection (NEW)
│ ├── renderers/
│ │ ├── graphviz_renderer.py # 2D Graphviz output
│ │ ├── plotly_renderer.py # 3D Plotly output
│ │ └── threejs_renderer.py # 3D Three.js output
│ └── utils/
│ ├── framework_detect.py # Auto-detection
│ └── grouping.py # Layer pattern grouping
├── tests/ # Test suite
├── examples/ # Demo scripts
├── docs/ # Documentation
└── pyproject.toml # Package config
# Run all tests
pytest tests/ -v
# With coverage
pytest tests/ --cov=modelviz --cov-report=html
# Run specific test
pytest tests/test_grouping.py -v- Branching graph support (ResNet, UNet skip connections)
- Transformer attention pattern visualization
- Interactive web dashboard
- Custom color themes
- Model comparison (side-by-side)
- FLOPs/MACs calculation
- ONNX model support
We welcome contributions! See CONTRIBUTING.md for guidelines.
git clone https://github.com/shreyanshjain05/modelviz.git
cd modelviz
python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev,torch,tf]"
pytest tests/ -v- Python 3.10+
- Type hints on all public functions
- Google-style docstrings
- Black + isort formatting
Apache 2.0 License - see LICENSE for details.
Made with ❤️ for the deep learning community
⭐ Star this repo if you find it useful!
☕ You can also support me on Ko-fi: https://ko-fi.com/shreyanshjain05 — every coffee keeps me going!