NeuroSchemaX

Neural network architecture visualization and export, powered by NN-SVG.

NeuroSchemaX parses neural-network models (ONNX, PyTorch, TensorFlow, or hand-written JSON/YAML specs), normalises them into a semantic representation, and renders them using the NN-SVG JavaScript engine — producing standalone HTML and SVG diagrams suitable for papers, theses, READMEs, and documentation.

Best-supported targets: MLP/FCNN and sequential CNN-style architectures (LeNet-style, AlexNet/VGG-style). ResNet and U-Net architectures render as honest labeled summaries (Residual Blocks, Encoder/Bottleneck/Decoder). Transformer architectures render as conceptual block summaries. Exact skip/concat paths, attention flow, and arbitrary DAGs are not drawn; they are preserved in export-debug-json.

NeuroSchemaX visualizes architecture diagrams and inferred metadata. It does not visualize learned weights, activation values, real feature-map contents, exact attention flow, or exact skip/concat paths unless those are explicitly supported for the specific architecture type.

Open NeuroSchemaX Colab Viewer

What it does

1. Parse — reads ONNX, PyTorch, Keras, JSON, or YAML and understands the layer structure.
2. Analyse — detects layer types, skip connections, and block groupings, then recommends the best diagram style.
3. Render — produces standalone offline HTML or SVG (via headless Chromium), plus JSON export formats.

Installation

Install from PyPI:

pip install neuroschemax

Install from GitHub:

pip install git+https://github.com/arashsajjadi/NeuroSchemaX.git

Optional extras:

pip install "neuroschemax[onnx]"   # ONNX input (already in base install)
pip install "neuroschemax[torch]"  # PyTorch model input
pip install "neuroschemax[tf]"     # TensorFlow / Keras model input
pip install "neuroschemax[svg]"    # SVG export via headless Chromium
playwright install chromium        # also required for SVG export

pip install "neuroschemax[colab]"  # Colab / Jupyter inline display (IPython)
pip install "neuroschemax[dev]"    # tests and linter
pip install "neuroschemax[all]"    # everything except Playwright/Chromium

PyTorch → ONNX (recommended for real PyTorch models):

pip install "neuroschemax[torch]" onnxscript

onnxscript is required by torch.onnx.export in torch >= 2.x.

Google Colab:

!pip install "neuroschemax[colab]"

Save and download HTML for full interactive rendering — Colab's inline rendering is limited for complex JS diagrams.

Quickstart

import neuroschemax as nsx

nsx.draw("model.onnx")
nsx.savefig("architecture.html")

neuroschemax draw model.onnx

Open the generated HTML in any browser — no internet connection required.

Python API

Simplified stateful API

import neuroschemax as nsx

nsx.draw("model.onnx")        # parse and stash
nsx.savefig("diagram.html")   # use stashed arch — format inferred from extension
nsx.save_html("out.html")     # also HTML
nsx.show()                    # open in browser (inline in Jupyter)

# Or pass source directly
nsx.save_html("out.html", "model.onnx")

Figure object API

fig = nsx.figure(width=1400, height=700, theme="paper")
fig.draw("model.onnx")
fig.savefig("diagram.html")
fig.save_html("diagram.html")
fig.save_svg("diagram.svg")       # needs Playwright
fig.show()
fig.export_debug_json("debug.json")
fig.export_paper_json("paper.json")
fig.export_nnsvg_json("spec.json")

# Matplotlib-style sizing
fig = nsx.figure(figsize=(12, 6), dpi=120, theme="paper")

# Chaining
nsx.figure().draw("model.onnx").save_html("out.html")

Explicit functional API

nsx.parse_model(source)               # SemanticArchitecture
nsx.summarize_model(source)           # str
nsx.recommend_view(source)            # dict: family/confidence/is_approximate/reason/warnings

nsx.build_nnsvg_spec(source, **kw)    # NNSVGSpec
nsx.render_network_html(source)       # HTML string
nsx.render_network_svg(source)        # SVG string (needs Playwright)
nsx.save_network_html(path, source)   # Path
nsx.save_network_svg(path, source)    # Path (needs Playwright)

nsx.export_paper_json(source)         # JSON string
nsx.export_debug_json(source)         # JSON string
nsx.save_paper_json(path, source)     # Path
nsx.save_debug_json(path, source)     # Path
nsx.save_nnsvg_spec(path, spec)       # Path

nsx.doctor()                          # dict with status/version/assets/deps/messages

Rendering keyword arguments

These are accepted by every rendering function and nsx.figure():

label_mode:

• auto — compact labels for small models; name-only for large models
• name — layer name only (shortest; never overlaps)
• compact — operation type + key parameters: Conv 64 k3, MaxP k2 s2, Dense 128, GAP
• shape — shape only, no name
• full — layer name + complete shape string; may be wide for large models

Activations following a layer are fused into the label when show_activations=True: Conv 64 k3 +ReLU, Dense 128 +GeLU. BatchNorm, LayerNorm, and Dropout appear as inline badges: +BN, +LN, +Drop 0.5.

detail_level:

• auto — all layers for small models; grouped blocks for models with more than 12 spec layers
• summary — groups conv/pool sequences into named blocks with metadata, e.g. Block 2 / 4×Conv k3, 128ch / Pool ↓2
• full — every individual layer shown; intended for inspection and debugging, not screenshots

transformer_mode:

• block_summary — approximate conceptual block sequence: Tokens/Input → Embedding → [MH-Attn] / Add & Norm → [FFN] / Add & Norm → [Head]. This is not exact Transformer rendering. Q/K/V projections, individual heads, exact residual paths, and tensor flow are not drawn.
• unsupported — renders a structured HTML diagnostic card instead of a diagram, listing detected components and suggesting block_summary or debug JSON

approximate_mode:

• warn — amber warning badge shown in HTML (default)
• error — raises RenderError before rendering an approximate diagram
• allow — suppresses warning badges

CLI

# Draw (output defaults to <model>.html)
neuroschemax draw model.onnx
neuroschemax draw model.onnx -o diagram.html --theme paper

# Render with full options
neuroschemax render model.onnx -o diagram.html --theme thesis --width 1600
neuroschemax render model.onnx -o diagram.svg

# Inspect and summarise
neuroschemax inspect        model.onnx
neuroschemax summarize      model.onnx
neuroschemax summarize      model.onnx --format markdown
neuroschemax recommend-view model.onnx

# JSON exports
neuroschemax export-paper-json model.onnx -o model.paper.json
neuroschemax export-debug-json model.onnx -o model.debug.json
neuroschemax export-nnsvg      model.onnx -o model.nnsvg.json

# Environment diagnostics
neuroschemax doctor

Supported inputs

Supported outputs

Diagram families and fidelity

NN-SVG supports three sequential diagram families. NeuroSchemaX selects the best fit automatically based on the model structure.

NeuroSchemaX does not claim to render arbitrary DAGs, exact ResNet skip edges, exact U-Net encoder-decoder skip topology, or exact Transformer attention flow. Models that fall outside the three sequential families are shown as honest approximate summaries with clear on-diagram markers (+skip collapsed, concat collapsed) and amber warning badges in the HTML.

ResNet summary: Stem → Residual Block N (n×Conv k3, Cch, +skip collapsed) → Head / Classifier

U-Net summary: Encoder → Bottleneck → Decoder (concat collapsed) → Segmentation Head

Transformer block summary: Tokens/Input → Embedding → [MH-Attn] / Add & Norm → [FFN] / Add & Norm → [Head]

The complete layer graph (every Add/Concat/Upsample with attributes) is preserved in export-debug-json regardless of what is shown visually.

How architecture recommendation works

info = nsx.recommend_view("model.onnx")
# {
#   "family": "alexnet",
#   "confidence": "high",
#   "is_approximate": False,
#   "reason": "Deep CNN with 8 convolutional layers, mapped to AlexNet",
#   "warnings": [],
#   "complexity_hint": "sequential"
# }

Selection rules (in priority order):

• Attention layers → block-level LeNet view (LOW confidence, is_approximate=True)
• Recurrent layers (LSTM/GRU) → block-level LeNet view (LOW confidence)
• No convolutions, at least one dense layer → FCNN (HIGH confidence)
• 1–3 conv layers → LeNet (HIGH; MEDIUM with skip/merge warning)
• 4+ conv layers → AlexNet (HIGH; MEDIUM if skip connections present)

Troubleshooting

Run neuroschemax doctor first.

SVG export fails:

pip install playwright
playwright install chromium

Missing assets:

pip install --force-reinstall neuroschemax

HTML is blank: Open in a modern browser (Chrome/Firefox/Edge). Check the browser console for JS errors. Try --theme debug for extra output.

See docs/troubleshooting.md for more.

FAQ

Q: Does it work offline? Yes. Generated HTML files are fully self-contained — all JS is embedded.

Q: Can I use it in a Jupyter notebook? Yes. If fig is the last expression in a cell, _repr_html_() renders the diagram inline automatically. fig.to_html() returns the HTML string for manual display() calls. nsx.show() / fig.show() open in the browser outside notebooks.

Q: Can I use it in Google Colab? Yes, with limitations. Install with !pip install "neuroschemax[colab]". Inline rendering in Colab is limited — use fig.save_html("diagram.html") and download the file for full interactivity.

Q: My PyTorch ONNX export fails with No module named 'onnxscript'. torch >= 2.x requires onnxscript for torch.onnx.export. Install it: pip install onnxscript.

Q: Diagram shapes show ?. Shape propagation is best-effort. Common causes: dynamic batch axes (use a concrete size at export), or BatchNorm/Dropout removed by ONNX exporter in eval mode (expected — they are folded out of the computation graph).

Q: Does it support ONNX opset X? It reads graph topology and layer attributes regardless of opset version. Unsupported op types are normalised to LayerKind.UNKNOWN.

Q: Can I customise the diagram further? Override any rendering option via kwargs. For pixel-level control, export the NN-SVG JSON spec and load it in the NN-SVG web app.

Q: Why is confidence MEDIUM/LOW? The model has features that don't map perfectly to the chosen NN-SVG family. The diagram is still generated; read the warnings for specifics.

Q: Why does the diagram show fewer layers than my model has? Large models are grouped into conv/pool blocks by default (detail_level="auto"). Pass detail_level="full" to see every individual layer. Note that full mode is intended for inspection; for screenshots, detail_level="summary" or the default produces more readable output.

Development

git clone https://github.com/arashsajjadi/NeuroSchemaX.git
cd NeuroSchemaX
pip install -e ".[dev]"
pytest tests/ -v
ruff check src/ tests/

See CONTRIBUTING.md.

Author

Arash Sajjadi — maintainer and author.

GitHub: arashsajjadi/NeuroSchemaX

License

MIT — see LICENSE.

NN-SVG is by Alex Lenail and also MIT-licensed. NeuroSchemaX integrates it as an embedded rendering engine.