johnmarktaylor91/torchlens

TorchLens – Deep Inspection of PyTorch Models

What it does – TorchLens is a Python library that records every tensor that flows through a PyTorch model during a forward (and optionally backward) pass. It builds a full computational‑graph DAG, attaches a huge amount of per‑operation metadata (shapes, types, FLOPs, timing, RNG state, etc.), and can visualise the graph as a PDF/HTML. The capture is “provably faithful”: the library re‑plays the captured DAG and checks that the outputs match the original run, flagging any mismatch.

Key features

• One‑line tracing: log = tl.trace(model, input) returns a TraceLog object containing activations, gradients and metadata.
• Rich selectors – choose which ops to save, halt, or intervene on using predicates like tl.func('relu'), tl.in_module('classifier'), or custom logical combinations.
• Forward + backward capture – gradients per op are available; you can limit gradient saving to specific layers.
• Receptive / projective field analysis – geometric queries that are cross‑checked with empirical gradient masks.
• Interventions – zero‑ablate, scale, add noise, swap tensors, etc., on‑the‑fly during capture.
• Visualization – auto‑generated PDFs, rolled/unrolled views, interactive HTML, Netron export.
• Dataset‑wide extraction – batch‑process a dataset and write activations to disk.
• Performance‑aware – selective capture can be faster than a raw forward; full capture is ~14× slower but still benchmarked.
• Works with any PyTorch model – static, dynamic control flow, recurrent loops, transformers, GNNs, diffusion, MoE, etc.

Typical workflow

import torch, torchvision.models as models, torchlens as tl
model = models.resnet18(weights='IMAGENET1K_V1').eval()
x = torch.rand(1,3,224,224)
log = tl.trace(model, x)          # capture everything
print(log.summary())             # table of layers, FLOPs, params
print(log['conv2d_1_1'].out.shape)  # access a specific activation
log.draw()                       # PDF of the graph

You can then intervene, compare runs, or export the trace.

Installation – Requires Graphviz for graph drawing. Install via pip install torchlens (CPU‑only core) and optional extras like torchlens[hf] for HuggingFace models.

Who it’s for – Researchers and engineers who need fine‑grained introspection of neural‑network internals, debugging of dynamic architectures, or systematic activation/gradient extraction for analysis, probing, or model‑based interventions.

Limitations – Backward capture only works on PyTorch; other backends get leaf‑level gradients via a second AD pass. Platform support is primarily Linux/CPU; macOS/Windows have limited CI checks.

---

Links

• Paper: https://www.nature.com/articles/s41598-023-40807-0
• Model Menagerie (catalog of 8.5k captured models): https://modelmenagerie.ai
• Tutorials, performance guide, and API docs are in the repo under docs/ and notebooks/.