TorchSpace

Spatial debugging for PyTorch — architecture, activations and gradients in one interactive 3D scene

_{▲ Live Demo — the full interactive viewer with two demo models, running entirely in your browser (nothing to install, no server)}

TorchSpace renders a PyTorch model as one interactive Three.js scene: X·Y is the architecture (layers, blocks, skip connections), +Z is forward activation magnitude and −Z is backward gradient magnitude, both on a log scale. Vanishing gradients, activation explosions, dead regions and forward/backward asymmetry become visible shapes instead of numbers you have to hunt for. Built as a non-invasive extension of torchview — no code copied, full torchview feature set preserved.

Everything is captured as summary statistics only (never raw tensors) and exported as a single self-contained HTML file: three.js is inlined and the statistics are embedded, so it works in Colab, Jupyter, VS Code, as an email attachment, and offline.

Install

pip install torchspace

Quick start

import torchspace

run = torchspace.trace(model, input_data=x)   # structure + instrumentation
out = model(x)                                # captured forward
run.capture_backward(criterion(out, y))       # captured backward
run.show(height=620)                          # interactive scene, inline in a notebook
run.export_html("model.torchspace.html")      # ...or one shareable file
run.detach()                                  # de-instrument the model

torchspace.view(model, input_size=(1, 3, 224, 224), device="meta")
# structure-only mode, no activation capture (works on meta device)

Tutorial: TorchSpace_Tutorial.ipynb walks through the full API — quick start, diagnosing a pathological network, the IR, meta-device inspection, torchview-parity graphviz export, module attributes, training-loop instrumentation, and export/sharing. Open it with the Colab badge at the top of this page.

What you get

• One scene for the whole story — architecture, per-layer activation and gradient magnitude, execution order and timing, all aligned in space.
• Rule-based diagnostics with evidence — vanishing/exploding gradients, dead regions, NaN/Inf, forward/backward asymmetry. Click a finding to jump to the layer; select a layer to filter the panel to its findings.
• Timeline replay — animate the actual forward (green) / backward (purple) execution order with flow pulses.
• Real training loops — statistics are captured by ordinary hooks on your actual training steps, on-device, detached immediately. No separate profiling pass, no RecorderTensor near your hot path.
• Notebook-native — run.show() renders the scene in a sandboxed srcdoc iframe that works in Colab's network-restricted output frames.

Viewer controls

Drag = rotate · wheel = zoom · right-drag = pan · hover = exact stats · click = inspect (+ diagnostics filter) · double-click = expand/collapse · 2D Ortho = torchview-style top-down view · ▶ Replay = forward/backward replay · toggles for tensor geometry, ±Z overlays, ops, rolling, labels. The log-scale ruler beside the model gives the quantitative 1e^x scale.

torchview parity

The trace captures at full detail (torchview's information is a strict subset of the IR), module/function attributes are collected by default (collect_attributes=False to disable), and torchview's native graphviz output stays available at any granularity:

run.export_dot("model.gv", depth=3, graph_dir="LR")   # DOT source
run.draw_graph(depth=2, roll=True, save_graph=True,
               filename="model", format="png")        # torchview pipeline

draw_graph/export_dot accept every torchview.draw_graph option (depth, hide_inner_tensors, hide_module_functions, roll, graph_dir, strict, ...) and reuse the exact inputs of the original trace; runtime capture hooks are suspended for the auxiliary pass.

What ends up in an export (privacy)

Exports contain statistics only — never raw tensors. Two edges to know before sharing: collect_attributes=True (default) embeds a repr-style summary of each module's public attributes, which for custom modules can include checkpoint paths or config values — pass collect_attributes=False for sensitive models; and for a 1-element tensor the stats (min=max=mean) necessarily reveal its exact value.

Repository layout

torchspace/
  _tv_compat.py    # single audited import surface over torchview
  tracer.py        # non-invasive structural tracer (extension points A & B)
  runtime.py       # forward/backward hook capture on plain tensors
  stats.py         # compact tensor statistics (never raw tensors)
  ir.py            # IR builder: stable IDs, empty-pass contraction, edges
  diagnostics.py   # rule-based anomaly detection
  api.py           # public view()/trace()/SpaceRun (+ draw_graph/export_dot)
  export.py        # embeds IR into the self-contained viewer
  assets/viewer.html       # packaged viewer (generated — do not edit)
viewer/
  src.html                 # viewer source (edit this)
  three_r128_inline.html   # vendored three.js r128 (script-wrapped)
scripts/build_viewer.py    # src.html + three -> torchspace/assets/viewer.html
docs/                      # GitHub Pages live demo + README assets
build_demo.py              # end-to-end driver -> out/torchspace_viewer.html
tests/test_smoke.py

Developing the viewer: edit viewer/src.html, then run python scripts/build_viewer.py to regenerate the packaged asset, and python build_demo.py to rebuild the demo scene.

Known limitations (PoC)

• Layout is a simple deterministic layered algorithm; production uses ELK compound layout. The 3D viewer has no graph_dir equivalent (use export_dot(graph_dir=...) for that).
• Tensor dtype/bytes are captured for op outputs and inputs, not for tensors created by torch.* creation ops inside forward.
• requires_grad in the IR reflects the (no-grad) structural pass.
• Single step / single device shown in the viewer; earlier steps stay in the IR. Static HTML export only (no live server).

Creator

License

MIT — torchview (MIT) is imported as a dependency, none of its code is copied; three.js r128 (MIT, © Three.js Authors) is bundled inside the packaged viewer with its license header retained.