johnmarktaylor91/torchlens

TorchLens – PyTorch 模型的深度检测

功能说明 – TorchLens 是一个 Python 库，可记录 PyTorch 模型在正向（以及可选的反向）传递过程中流经的每一个张量。它会构建完整的计算图 DAG，附加大量的算子元数据（形状、类型、FLOPs、时间、RNG 状态等），并能将计算图可视化为 PDF/HTML。此捕获过程是“可证明忠实的”：库会重播捕获的 DAG 并检查输出是否与原始运行结果相符，若有不符则会标记。

主要功能

• 单行追踪：log = tl.trace(model, input) 会返回一个包含激活值、梯度和元数据的 TraceLog 对象。
• 丰富的选择器 – 使用如 tl.func('relu')、tl.in_module('classifier') 或自定义逻辑组合等谓词，选择要保存、暂停或干预的算子。
• 正向与反向捕获 – 提供每个算子的梯度；您可以将梯度保存限制在特定层。
• 感受野 / 投射野分析 – 通过经验梯度掩码进行交叉验证的几何查询。
• 干预 – 在捕获过程中即时进行零消融 (zero-ablate)、缩放、添加噪声、交换张量等操作。
• 可视化 – 自动生成的 PDF、折叠/展开视图、交互式 HTML、Netron 导出。
• 数据集级提取 – 批处理数据集并将激活值写入磁盘。
• 性能感知 – 选择性捕获可能比原始正向传递更快；完整捕获约慢 14 倍，但仍经过基准测试。
• 适用于任何 PyTorch 模型 – 静态、动态控制流、循环、transformers、GNNs、扩散模型、MoE 等。

典型工作流程

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

您可以随后进行干预、比较运行结果或导出追踪记录。

安装 – 绘制计算图需要 Graphviz。通过 pip install torchlens 安装（仅 CPU 核心），并可选安装如 torchlens[hf] 等额外包以支持 HuggingFace 模型。

适用对象 – 需要对神经网络内部进行细粒度检测、调试动态架构，或进行系统性激活值/梯度提取以进行分析、探测或基于模型的干预的研究人员与工程师。

限制 – 反向捕获仅适用于 PyTorch；其他后端通过第二次 AD 传递获取叶节点梯度。平台支持主要为 Linux/CPU；macOS/Windows 的 CI 检查有限。

---

链接

• 论文：https://www.nature.com/articles/s41598-023-40807-0
• Model Menagerie (8.5k 个已捕获模型目录)：https://modelmenagerie.ai
• 教程、性能指南和 API 文档位于仓库的 docs/ 和 notebooks/ 文件夹中。