johnmarktaylor91/torchlens

TorchLens – What It Is

TorchLens is a Python library that lets you inspect, record, and modify what happens inside any PyTorch model (and, in preview mode, a few other deep‑learning frameworks). It works by wrapping every low‑level operation during a forward (and optional backward) pass, building a full directed‑acyclic graph (DAG) of the computation. From that graph you can:

• Save every activation and gradient (or just a filtered subset).
• Query activations by operation name, module path, or ordinal index.
• Visualize the whole graph as a PDF, with options for rolled (compact) or unrolled (each loop iteration shown) views.
• Compute rich metadata – shapes, dtypes, devices, FLOPs, timing, parameter counts, call‑stack location, RNG state, etc. (180+ fields per op, >550 fields overall).
• Intervene on the fly – zero‑ablate, scale, add noise, replace tensors, etc., during the forward pass.
• Replay a captured graph later, optionally with different interventions, enabling “what‑if” experiments without re‑running the original model code.
• Capture gradients per operation, so you can inspect backward‑pass information just as easily as forward activations.

The library also ships with a Model Menagerie – a publicly browsable collection of >11 600 neural‑network architectures (CNNs, transformers, GNNs, diffusion models, etc.) that have been automatically captured and verified with TorchLens.

---

Why It Matters

• Dynamic models – Because TorchLens hooks the eager Python execution rather than tracing a static graph, it correctly handles conditional branches, loops, recurrent unrolls, and fused ops that many other feature‑extraction tools miss.
• Provable fidelity – For the menagerie entries, TorchLens re‑plays the captured DAG and checks that outputs match the original forward pass, plus a suite of metadata‑invariant “tripwires”. About 89 % of the catalog is algorithmically verified, giving confidence that the recorded activations are not just plausible but provably correct.
• Fine‑grained control – Selectors (tl.func('relu'), tl.in_module('classifier'), logical combinations) let you record only the tensors you need, dramatically reducing post‑processing effort.
• Intervention API – Researchers can experiment with activation‑patching, ablation, or scaling without rewriting model code, which is useful for mechanistic interpretability, robustness testing, or debugging.
• Cross‑framework preview – Although the full feature set is currently PyTorch‑only, the library provides early support for JAX, tinygrad, MLX, Paddle, and TensorFlow, hinting at a future “one‑API‑for‑all‑DL‑frameworks”.

---

Core Concepts & API (as described in the README)

Concept	How You Use It	What You Get
Trace a model	log = tl.trace(model, input_tensor)	log – a TraceLog object containing the full DAG and all captured data
Select what to save	save=tl.func('relu'), save=tl.in_module('classifier'), logical combos (&, `	`)
Query activations	log['relu_1_2'].out, log[7].func_name	Direct access to tensors, shapes, gradients, etc.
Visualize	log.draw() (options: vis_mode='rolled' / 'unrolled')	PDF of the computational graph, with nesting and recurrence handling
Gradient capture	log = tl.trace(model, x, capture=tl.options.CaptureOptions(save_grads=True)) then log.log_backward(loss)	Per‑op gradient tensors available via log['op_name'].grad
Intervention	tl.when(tl.func('relu'), tl.zero_ablate()) passed via intervene=	The specified ops are modified (e.g., zeroed) during the forward pass
Replay / fork	log.fork(), log.replay(), log.rerun(model, x)	Run the captured graph again, optionally with different interventions
Metadata export	df = log.to_pandas()	Pandas DataFrames with all recorded fields for ops, modules, parameters
Distributed capture	tl.distributed.arm() before tracing	Rank‑local collectives become first‑class nodes; later merged with tl.merge_ranks

---

Typical Workflow (from README)

1. Install (requires Graphviz for drawing):

   sudo apt install graphviz
   pip install torchlens
   

2. Trace a model:

   import torch, torchvision.models as models, torchlens as tl
   model = models.alexnet(weights=None)
   x = torch.randn(1,3,224,224)
   log = tl.trace(model, x)
   

3. Inspect:

   print(log.summary())               # table of modules, FLOPs, etc.
   print(log['relu_1_2'].out.shape)   # activation shape
   log.draw()                          # PDF graph
   

4. Filter & intervene (e.g., zero‑ablate all ReLUs):

   ablated = tl.trace(model, x,
                      save=tl.func('relu'),
                      intervene=tl.when(tl.func('relu'), tl.zero_ablate()))
   

5. Analyze gradients (optional backward capture):

   x.requires_grad_()
   log = tl.trace(model, x, capture=tl.options.CaptureOptions(save_grads=True))
   loss = log[log.output_layers[0]].out.sum()
   log.log_backward(loss)
   print(log['relu_1_2'].grad.shape)
   

6. Export metadata for downstream analysis:

   df = log.to_pandas()
   df.to_csv('activations.csv')
   

---

Who Might Use TorchLens?

• Research scientists probing the internals of large vision, language, or multimodal models.
• Interpretability engineers who need precise activation‑patching or receptive‑field calculations.
• Performance engineers wanting to verify FLOPs, timing, or memory footprints of custom architectures.
• Educators demonstrating how data flows through dynamic or recurrent networks.
• Tool builders who need a reliable way to capture a model’s graph for downstream tooling (e.g., pruning, quantization, or conversion pipelines).

---

Limitations Mentioned

• Full‑feature support is currently PyTorch‑only; other frameworks are in preview and lack many capabilities (e.g., interventions, full backward capture).
• Capturing all activations can be 14× slower than a raw forward pass; selective saving or early halting can mitigate this.
• Very large models or those with exotic dependencies may remain unverified in the Model Menagerie.

---

Quick Reference Links (from README)

• Paper – https://www.nature.com/articles/s41598-023-40807-0
• Tutorial notebooks – notebooks/torchlens_in_10_minutes.ipynb, facets_tutorial.ipynb
• Performance guide – docs/performance.md
• Receptive/projective fields – docs/receptive_projective_fields.md
• AI‑agent reference – docs/for-ai-agents.md
• Limitations & remedies – docs/reference/limitations.md
• Backend overview – docs/backends.md

---

In short: TorchLens is a comprehensive, provenance‑focused toolkit for extracting, visualizing, and manipulating the full computational graph of PyTorch models (with preview support for other frameworks). It emphasizes correctness (graph replay verification), rich per‑operation metadata, and a flexible intervention API, making it valuable for deep‑learning research and engineering tasks that require deep insight into model internals.