TorchLens

See, save, and steer any PyTorch model. TorchLens captures every activation and gradient -- across the forward and backward pass -- auto-visualizes the full computational graph, exposes rich per-op metadata, and lets you intervene on the network as it runs. Any architecture, even dynamic and recurrent ones.

Explore the Model Menagerie -- a live, browsable atlas of 11,600+ cataloged neural-network architecture entries captured with TorchLens, from McCulloch & Pitts (1943) to today's frontier models. (Early preview.)

Run across 11,600+ cataloged entries in the Model Menagerie (image, video, audio, multimodal, language; feedforward, recurrent, transformer, GNN, MoE, diffusion). Roughly 89% are currently algorithmically verified for capture correctness: each verified capture is replayed op-by-op against its own forward pass, with metadata-invariant tripwires over the graph, so faithful capture is proven, not assumed. TorchLens also records 180+ metadata fields per operation, and 550+ fields in total across every record type — operations, modules, parameters, buffers, gradients, and the model itself.

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)     # one call -- full graph + all activations
print(log.summary())          # module table, op count, FLOPs
print(log['relu_1_2'].out.shape)   # grab any activation by name ...
print(log['features.6'].out.shape) # ... or by module path
print(log[7].func_name)            # ... or by ordinal
log.draw()                    # PDF of the computational graph

Quick Links

• Paper | 10-minute tutorial notebook | Facets tutorial | 5-minute gallery | 50-minute gallery
• Performance guide | Receptive & projective fields | AI-agent quick reference | Limitations and remedies | Migration tables

Validated across 11,600+ catalog entries

TorchLens is not just smoke-tested on example models. Its menagerie validation campaign runs the same adversarial check across more than 11,600 cataloged architecture entries: capture the model with TorchLens, forward-replay the captured DAG, compare replayed outputs against the original forward pass, and run metadata-invariant tripwires over the resulting graph. If replay or an invariant fails, the capture is treated as genuinely wrong and caught automatically, not waved through because the model "ran."

model forward -> TorchLens capture -> DAG replay -> output parity + metadata invariants

Today, roughly 89% of that 11,600+ catalog is algorithmically verified and climbing, covering about 5,400 distinct architecture families after variants are collapsed. That is the wedge: TorchLens aims for captures that are provably faithful, not just plausible. Plain forward hooks and static extraction utilities can be fast and useful, but they can also silently miss dynamic paths, reused modules, functional ops, fused attention, recurrent unrolls, or metadata needed to tell two sites apart.

The residual is tracked openly: the remaining tail is mostly genuinely huge models, models with hard-to-build dependencies, or architectures that need more bespoke inputs before an automated run is fair. They are counted as unverified, not hidden as successes.

Installation

Install Graphviz first (required for graph visualizations), then TorchLens:

sudo apt install graphviz   # Debian/Ubuntu; see graphviz.org for other platforms
pip install torchlens

Compatible with PyTorch 2.1+.

Quickstart

import torch
import torchvision.models as models
import torchlens as tl

model = models.alexnet(weights=None)
x = torch.randn(1, 3, 224, 224)

log = tl.trace(model, x)
print(log.summary())

Model: AlexNet
+-----------------------------+---------------+--------+-------+
| Layer                       | Output Shape  | Params | Train |
+-----------------------------+---------------+--------+-------+
| input                       | [1,3,224,224] | 0      | -     |
| features (Sequential)       | [1,256,6,6]   | 2.5 M  | yes   |
| avgpool (AdaptiveAvgPool2d) | [1,256,6,6]   | 0      | -     |
| classifier (Sequential)     | [1,1000]      | 58.6 M | yes   |
| output                      | [1,1000]      | -      | -     |
+-----------------------------+---------------+--------+-------+
Params: 61,100,840 unique; trainable: 61,100,840
Ops: 22 total
Edges: 23 total
Forward FLOPs: 1.4 GFLOPs  MACs: 718.9 MFLOPs

Index any operation by name, module path, or ordinal:

log['relu_1_2'].out.shape      # torch.Size([1, 64, 55, 55])
log['features.6'].out.shape    # same op via module path
log[7].func_name               # 'conv2d'
log['conv2d_3'].out.shape      # short name (ordinal suffix optional)
log[-1].layer_label            # 'output_1'

Visualize the graph as a PDF:

log.draw()                        # unrolled by default
log.draw(vis_mode='rolled')       # rolled (compact for recurrent)
log.draw(vis_mode='unrolled')     # every pass as a distinct node

What You Can Do

1. Flexible feature extraction

Save everything, or select exactly what you need:

# Save only relu activations
log = tl.trace(model, x, save=tl.func('relu'))

# Save all ops inside the 'classifier' submodule
log = tl.trace(model, x, save=tl.in_module('classifier'))

# Save conv2d ops that are immediately followed by a relu, keeping a 4-op lookback window
conv_before_relu = tl.func('conv2d') & tl.followed_by(tl.func('relu'))
log = tl.trace(model, x, save=conv_before_relu,
               lookback=4, lookback_payload_policy='detached_raw')

# Stop capture early (can be faster than a plain forward pass)
log = tl.trace(model, x, save=tl.in_module('features.6'), halt=tl.in_module('features.6'))

# Lightweight sparse recording for tight loops -- materialize structure later
recording = tl.record(model, x, save=tl.func('relu'))
trace = recording.to_trace()

# One-line activation pull
act = tl.pluck(model, x, 'relu_1_2')   # returns tensor directly

# Batch extraction across a dataset (any iterable of unbatched samples works)
dataset = [torch.randn(3, 224, 224) for _ in range(8)]
tl.extract_dataset(model, dataset, layers=['relu_1_2', 'conv2d_3_7'],
                   batch_size=32, output_dir='torchlens-activations/')

Performance note: With halt= and tl.record, capture can run faster than the raw forward pass -- measured at 0.84x raw on ResNet-18 and 0.83x on GPT-2 (HookedTransformer) at 25% depth. Full exhaustive capture runs at roughly 14x the raw forward and amortizes on large models. See docs/performance.md for the full benchmark table.

Save and load traces portably:

tl.save(log, 'my_trace')
loaded = tl.load('my_trace')

2. Forward AND backward pass

Capture per-op gradients with the same API:

x = torch.randn(1, 3, 224, 224, requires_grad=True)
log = tl.trace(model, x, capture=tl.options.CaptureOptions(save_grads=True))
log.log_backward(log[log.output_layers[0]].out.sum())

grad = log['relu_1_2'].grad      # gradient tensor flowing through that op
print(grad.shape)                 # torch.Size([1, 64, 55, 55])

Narrow gradient saving to specific ops with the same selector predicates:

log = tl.trace(model, x, capture=tl.options.CaptureOptions(save_grads=tl.func('relu')))
log.log_backward(log[log.output_layers[0]].out.sum())

Backward capture is PyTorch-only. Non-torch backends expose derived leaf-level gradients through a second AD pass. See docs/backward.md.

Influence geometry in both directions

The first tool that gets ResNets right: TorchLens computes receptive fields back toward an input and projective fields forward toward an output over the captured DAG, then can cross-check the geometric answer with an empirical gradient support mask. It includes ResNet skip connections because it analyzes executed graph paths rather than multiplying a module list.

target = log["features.6"]
rf = target.receptive_field
unit = rf.center_unit(batch_index=0)
box = rf.at((3, 3))
check = rf.check(unit)
# Projective geometry needs a windowed path; AlexNet's dense classifier head is not,
# so select a downstream conv endpoint explicitly with target=.
outgoing = target.projective_field.at((3, 3), target=log['features.8'])

See the receptive and projective fields guide for the status contract, visual overlays, validation, and layer-to-layer queries.

3. Vast metadata per operation

Every operation records shape, dtype, device, timing, FLOPs, parameter info, module containment, graph distances, conditional context, RNG state, and more. The full print of any op includes all of this:

print(log['conv2d_3_7'])

Layer conv2d_3_7, operation 7/22:
    Output tensor: shape=(1, 384, 13, 13), dtype=torch.float32, size=253.5 KB
        tensor([[-0.0198,  0.0946,  0.1109, ...
    Related Layers:
        - parent layers: maxpool2d_2_6
        - child layers: relu_3_8
    Params: Computed from params with shape (384, 192, 3, 3), (384,); 663936 params total (2.5 MB)
    Function: conv2d (grad_fn_handle: ConvolutionBackward0)
    Computed inside module: features.6:1
    Config: out_channels=384, in_channels=192, kernel_size=(3, 3), padding=(1, 1)
    Time elapsed: 1.4 ms
    Lookup keys: -17, 7, conv2d_3, conv2d_3:1, conv2d_3_7, conv2d_3_7:1, features.6, features.6:1

Every op also records the Python call stack that produced it, with file and line number:

loc = log['conv2d_3_7'].code_context[0]
print(loc.file, loc.line_number, loc.func_name)

Metadata is available as pandas DataFrames:

df = log.to_pandas()            # one row per op
params_df = log.params.to_pandas()
modules_df = log.modules.to_pandas()

4. Automatic visualization

log.draw()                           # default: unrolled with sibling ordering
log.draw(vis_mode='rolled')          # compact rolled layout
log.draw(vis_mode='unrolled')        # every pass as a distinct node

Control nesting depth to zoom in on submodules:

For recurrent models, the rolled view collapses repeated structure cleanly:

class SimpleRecurrent(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(in_features=5, out_features=5)
    def forward(self, x):
        for r in range(4):
            x = self.fc(x)
            x = x + 1
            x = x * 2
        return x

recurrent_model = SimpleRecurrent()
seq = torch.randn(6, 5)
recurrent_log = tl.trace(recurrent_model, seq)
print(recurrent_log['linear_1:2'].out)     # second pass of the linear layer
recurrent_log.draw(vis_mode='rolled')

5. Interventions

Ablate, steer, scale, or replace activations during the forward pass:

# Zero-ablate all relu activations inline during capture
ablated = tl.trace(model, x, save=tl.func('relu'),
                   intervene=tl.when(tl.func('relu'), tl.zero_ablate()))
print(ablated['relu_1_2'].out.abs().max())  # tensor(0.)

# Scale relus to 50%
scaled = tl.trace(model, x, save=tl.func('relu'),
                  intervene=tl.when(tl.func('relu'), tl.scale(0.5)))

Available helpers: tl.zero_ablate, tl.mean_ablate, tl.resample_ablate, tl.steer, tl.scale, tl.clamp, tl.noise, tl.project_onto, tl.project_off, tl.swap_with, tl.splice_module.

For post-hoc DAG replay and isolated experiments, capture with intervention_ready=True and use log.fork() + log.replay() / log.rerun(model, x). Live hooks during rerun require capture-time selectors (e.g. tl.func(...), tl.module(...)); finalized labels resolve via log.find_sites(...). See docs/intervention_api.md for the full reference.

Compare multiple runs side by side with tl.bundle:

clean_log = tl.trace(model, x, save=tl.func('relu'))
patched_log = tl.trace(model, x, save=tl.func('relu'),
                       intervene=tl.when(tl.func('relu'), tl.zero_ablate()))
bundle = tl.bundle({'clean': clean_log, 'patched': patched_log}, baseline='clean')
bundle.compare_at('relu_1_2')   # one site; a multi-site selector must resolve uniquely

Facets provide named sub-views for attention heads, LSTM outputs, and fused projections (for models with those structures):

The following is an API sketch; vit_model and lstm must be models whose module structures provide the shown paths, and are not defined by this generic example.

# API sketch; `vit_model` and `lstm` are application-supplied models with these paths.
# ViT / transformer model with attention blocks
log = tl.trace(vit_model, x)
q = log.modules['blocks.0.attn'].facets['q']    # query vectors for head 0
h_n = log.modules['lstm'].facets['h_n']         # LSTM final hidden state

See docs/facets.md for the full facets reference, including activation patching helpers, SDPA reconstruction, and TransformerLens aliases.

See docs/intervention_api.md for the full selector and helper reference.

6. Works on anything, including dynamic and recurrent models

TorchLens uses eager-mode Python-level function wrapping rather than graph tracing. This means it captures whatever actually runs, including:

• Dynamic control flow (if/else branching, loops, early exits)
• Recurrent architectures (RNNs, LSTMs, state-space models)
• Transformer variants including fused attention
• Graph neural networks
• Mixed architectures

This is the key differentiator from static-graph extractors like torchvision.feature_extraction, which require static computational graphs and cannot handle dynamic architectures.

Distributed boundaries. With tl.distributed.arm() enabled before rank-local capture, explicit in-forward torch.distributed Python collectives become first-class boundary nodes. Diagnose rank sets with tl.merge_report(...) and merge compatible rank traces with tl.merge_ranks(...). Sharded tensor topologies such as DTensor/FSDP/TP and pipeline point-to-point graphs still refuse with typed findings; see the merged-trace contract.

Multi-backend. The same tl.trace API works across frameworks via backend=:

Preview save= selectors filter what is exposed, not what is captured (no memory reduction); preview tl.validate(...) returns a status whose bool() raises for partial coverage; and every preview auto-routes genuine framework models on backend=None. See docs/backends.md for the per-backend contract.

# API sketch; supply compatible models and input in an application context.
log = tl.trace(torch_model, x)                      # PyTorch (default)
log = tl.trace(jax_fn,      inputs, backend='jax')  # JAX preview
log = tl.trace(tg_fn,       inputs, backend='tinygrad')
log = tl.trace(paddle_model, x,     backend='paddle')
log = tl.trace(tf_model,    x,      backend='tf')

PyTorch remains the full-feature backend. Preview backends are pinned and documented in docs/.

Gallery

TorchLens visualizes any architecture -- no matter how exotic. Explore the Model Menagerie: a browsable atlas of 11,600+ cataloged neural-network architecture entries -- from McCulloch & Pitts (1943) to today's frontier models -- each with structured metadata and a TorchLens-rendered diagram. Roughly 89% currently carry the replay-and-invariant verification described above.

Early preview. The gallery is live and growing; full-text search, a downloadable dataset, and richer per-model pages are on the way.

A sample across families is shown below.

Classic CNN + Vision Transformer

State-Space + Recurrence

Mixture-of-Experts + Generative

Graph Networks + Exotic

Reinforcement Learning + Quantum ML + Scale

Pin your architecture in CI

Use the provisional address-free structural hash to catch an unintended graph change without pinning model weights or module names:

# `model` and `x` are your model and a representative example input.
import torchlens as tl

pinned = tl.assert_unchanged(model, x, expected=None)  # prints and returns a hash
tl.assert_unchanged(model, x, pinned)  # raises if the architecture changes

See tl.hash for trace-level hashing and its structural scope.

Compatibility

Before filing a bug for a model-specific failure, run the runtime compatibility report:

compat = tl.compat.report(model, x)
print(compat.to_markdown())

tl.compat.report inspects the model wrapper, modules, parameter sharing, input tensors, CUDA visibility, and common framework markers, then reports each row as pass, known_broken, scope, or not_tested.

torch.compile coexists with capture. On torch >= 2.6, every capture holds the public torch.compiler.set_stance("force_eager") scoped to the forward, so compiled regions run their original eager Python: the interior is fully logged with full verified semantics, zero graph breaks or new compiles happen during capture, compiled caches stay intact, and wrapper install/uninstall costs at most one bounded recompile on the next compiled call. On torch < 2.6 the historical fallback holds: a Dynamo-traced region reached mid-capture is bypassed with a one-per-forward warning and the returned trace honestly contains only what ran outside it (capture_verified=False, reason "dynamo_region_not_logged"). TorchLens remains not compatible with TorchScript or torch.export -- those forwards do not run as ordinary Python, so the wrappers cannot intercept ops. It also has specific behaviors around FSDP, sparse tensors, meta tensors, quantization, and torch.func.vmap.

See LIMITATIONS.md for the full matrix: what fails, what works, and the recommended workaround for each context.

TorchLens recovers most detached from torch import ... references with a disclosed rescue re-run and a small mechanical belt. The historical broad sys.modules crawl and patch_policy= rollout are deleted; those arguments are deprecated no-ops. For the strongest and simplest guarantee, call torchlens.backends.torch.wrappers.wrap_torch() before creating detached references. The optional escape_detector="shadow" diagnoses raw callable escapes. See detached-reference handling and the limitations catalog.

Tutorials and Docs

Security

Portable bundles contain a pickle file in metadata.pkl. Only load bundles from trusted sources. Loading an untrusted bundle with tl.load() can execute arbitrary code.

Other Packages You Should Check Out

TorchLens focuses on activation extraction, graph visualization, and intervention and intentionally omits model loading, stimulus management, and analysis pipelines. These packages cover that ground well:

• Cerbrec: interactive visualization and debugging for deep neural networks (uses TorchLens under the hood for PyTorch graph extraction)
• ThingsVision: model loading, stimulus management, and representational analysis for vision models
• Net2Brain: end-to-end pipeline for comparing DNN representations to neural data
• surgeon-pytorch: lightweight activation extraction with training-loss hooks
• deepdive: model loading and benchmarking across many model families
• torchvision feature_extraction: fast activation extraction for models with static computational graphs
• rsatoolbox: representational similarity analysis for DNN activations and brain data

Acknowledgments

The development of TorchLens benefitted greatly from discussions with Nikolaus Kriegeskorte, George Alvarez, Alfredo Canziani, Tal Golan, and the Visual Inference Lab at Columbia University. Thank you to Kale Kundert for helpful discussion and code contributions enabling PyTorch Lightning compatibility. Network visualizations are generated with Graphviz. Logo created by Nikolaus Kriegeskorte.

Citing TorchLens

To cite TorchLens, please cite this paper:

Taylor, J., Kriegeskorte, N. Extracting and visualizing hidden activations and computational graphs of PyTorch models with TorchLens. Sci Rep 13, 14375 (2023). https://doi.org/10.1038/s41598-023-40807-0

If you find TorchLens useful, a star on this repo is appreciated.

Contact

TorchLens is in active development. Questions, bug reports, and suggestions are welcome via email, Twitter, the issues page, or the discussion board.