TorchLens for debugging and training networks¶

When a network misbehaves, the question is usually "which operation first went wrong, and what fed it": the first NaN, the layer whose activations drifted, the branch that ran when it should not have, the tensor that is eating memory. TorchLens records every operation of a forward pass with its inputs, outputs, shapes, bytes, source line and module, so these questions become lookups on a trace instead of a hook-and-print session, and the same records can be sampled during training to watch distributions drift over time.

The lists below start with the one-off questions and move to monitoring a whole training run. Every link is to a file in this repository.

Start here¶

• Tutorial notebooks: eleven short notebooks on laptop-sized models. Notebook 01 captures a forward, 02 reads a number out, 03 draws the graph, 04 keeps memory in check and 06 reads gradients.
• Find the first NaN notebook: the first NaN- or Inf-producing operation, in two lines.
• See the graph and inspect any operation: a picture of every operation the forward ran, and the output, module and arguments of any one of them, including operations no module returns.
• Model summaries: a one-call, metadata-only pass that lists every layer with shapes, parameter counts and cost.
• Debug utilities example script: a tiny model with one unstable operation and the calls that find it.

Locate a fault¶

• tl.debug reference: locate non-finite saved tensors, walk dependencies, compare runs, rank cost, when the graph itself is the diagnostic.
• Extracting source code with TorchLens: the call stack and the actual source line of every operation.
• Conditional branch attribution: which Python if branch each operation ran under, so you can see the control flow the forward actually took.
• Containers: the input and output container structure of each operation, for faults inside tuples, dicts and lists.
• Buffers: BatchNorm running statistics and other registered buffers, recorded alongside activations.
• Compare two runs in a bundle: the same model before and after a change, diffed operation by operation.

Gradients and the backward pass¶

• Gradients tutorial notebook: what drives one unit, read from the trace.
• Backward capture and the backward tutorial notebook: per-layer output gradients, the autograd graph connected back to forward layers, and validation of the backward record.
• The gradient contract: what TorchLens promises about saved gradients.

Memory and speed¶

• Memory debugging: which question, which tool: which tool answers which memory question, and where TorchLens's per-operation byte accounting fits beside the PyTorch allocator tools.
• Keep what you need: save only the activations you want when memory or time runs out.
• Performance guide and choosing a fast path: the capture modes and what each costs.
• Profilers to use alongside TorchLens: the PyTorch profiler, Perfetto and friends measure what TorchLens does not.
• Native PyTorch tools: what each is authoritative for: when the built-in tool is the right answer.

Monitor training¶

• Monitor training: watch how the distribution at every site drifts over training, find where it broke, and keep the history as a portable, mergeable artifact.
• A percentile-band training ribbon: the ribbon plot of one site's distribution over steps, from watch histograms.
• Training sanity checks (the checks kit): range probes and sessions that flag the usual training faults.
• Trackers: TensorBoard and wandb: feed TorchLens statistics to the tracker you already log to.
• Training on saved activations: when a saved activation is an input to a loss that will feed backward().
• Train a truncated backbone through trace.run: cut a pretrained network at a stage and fine-tune the prefix, from one trace.
• Weight spectra with WeightWatcher: per-layer power-law metrics of the weight matrices.
• Streaming statistics: dataset-level statistics of activations or gradients in constant memory.

Save, share and export¶

• Save and load a trace notebook: round-trip a trace as a .tlspec bundle.
• tl.export reference: static graph files, performance-viewer formats, tabular files and tracker objects.
• The Netron export and Google Model Explorer export: open the captured graph in a viewer.
• The single-file offline HTML report: one file to send to a colleague.
• Change a model's skeleton, keep your analyses aligned: what to do with saved analyses after you swap or wrap a block.

Connectors and migration¶

• Connectors to other tools: the PyTorch profiler, Chrome tracing and Perfetto, TensorBoard, Weights and Biases, MLflow and Aim, Netron, treescope, lovely-tensors and the rest, one entry each.
• Coming from torchinfo, torchsnooper, torchview or plain forward hooks.
• Glossary: the definition of every public name used above.