TorchLens

快速链接

• 介绍TorchLens的论文
• Colab教程
• 模型可视化“动物园”
• TorchLens提供的元数据

概述

TorchLens 是一个用于完成以下两件事的工具包：

1. 无需任何修改，只需一行代码即可轻松提取 PyTorch 模型中每一个中间操作的激活值。“每一个操作”指的就是每一个操作；“一行”就是真正的一行。
2. 通过直观的自动可视化以及关于网络计算图的丰富元数据（部分列表在此），帮助理解模型的计算结构。

下面是一个非常简单的循环模型的示例：正如你所看到的，你只需像平常一样定义模型并将其传入，TorchLens 就会返回完整的前向传播日志以及可视化结果：

class SimpleRecurrent(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = 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


simple_recurrent = SimpleRecurrent()
model_history = tl.log_forward_pass(simple_recurrent, x,
                                    layers_to_save='all',
                                    vis_mode='rolled')
print(model_history['linear_1_1:2'].activation)  # 第一次线性层的第二次传递

'''
tensor([[-0.0690, -1.3957, -0.3231, -0.1980,  0.7197],
        [-0.1083, -1.5051, -0.2570, -0.2024,  0.8248],
        [ 0.1031, -1.4315, -0.5999, -0.4017,  0.7580],
        [-0.0396, -1.3813, -0.3523, -0.2008,  0.6654],
        [ 0.0980, -1.4073, -0.5934, -0.3866,  0.7371],
        [-0.1106, -1.2909, -0.3393, -0.2439,  0.7345]])
'''

再来看一个非常复杂的 Transformer 模型（swin_v2_b），其前向传播过程中有 1932 个操作；你可以获取到每一个操作的保存输出：

TorchLens 的目标是能够对任何 PyTorch 模型实现上述功能。你可以在这个 模型动物园 中看到许多模型可视化的示例。

安装

要安装 TorchLens，首先确保已安装 graphviz（生成网络可视化所需的工具），然后使用 pip 安装 TorchLens：

sudo apt install graphviz
pip install torchlens

TorchLens 兼容 PyTorch 1.8.0 及以上版本。

使用指南

以下是使用方法的简短演示；如需交互式演示，请参阅 CoLab 演示。

TorchLens 的主要函数是 log_forward_pass：当它被调用时，会对给定的模型和输入执行一次前向传播，并返回一个 ModelHistory 对象，其中包含中间层的激活值及配套的元数据，同时还会生成前向传播过程中每个操作的可视化表示：

import torch
import torchvision
import torchlens as tl

alexnet = torchvision.models.alexnet()
x = torch.rand(1, 3, 224, 224)
model_history = tl.log_forward_pass(alexnet, x, layers_to_save='all', vis_mode='unrolled')
print(model_history)

'''
AlexNet前向传播日志：
	模型结构：纯前馈，无分支；共23个模块。
	前向传播中计算了24个张量（4.8 MB）；保存了24个张量（4.8 MB）。
	参数操作16次（总参数数61100840个；248.7 MB）。
	随机种子：3210097511
	耗时：0.288秒
	模块层级：
		features:
		    features.0, features.1, features.2, features.3, features.4, features.5, features.6, features.7,
		    features.8, features.9, features.10, features.11, features.12
		avgpool
		classifier:
		    classifier.0, classifier.1, classifier.2, classifier.3, classifier.4, classifier.5, classifier.6
	层：
		0: input_1_0
		1: conv2d_1_1
		2: relu_1_2
		3: maxpool2d_1_3
		4: conv2d_2_4
		5: relu_2_5
		6: maxpool2d_2_6
		7: conv2d_3_7
		8: relu_3_8
		9: conv2d_4_9
		10: relu_4_10
		11: conv2d_5_11
		12: relu_5_12
		13: maxpool2d_3_13
		14: adaptiveavgpool2d_1_14
		15: flatten_1_15
		16: dropout_1_16
		17: linear_1_17
		18: relu_6_18
		19: dropout_2_19
		20: linear_2_20
		21: relu_7_21
		22: linear_3_22
		23: output_1_23
'''

你可以通过以下任意一种等效方式索引 ModelHistory 对象，来提取特定层的信息，包括其激活值和有用的元数据：

1. 层的名称（约定为 'conv2d_3_7' 表示第3个卷积层，整体的第7层）
2. 该层所属模块的名称（例如 'features' 或 'classifier.3'）
3. 层的序号位置（例如 2 表示第2层，-5 表示倒数第5层；这里的输入和输出也被计为层）。

为了快速确定这些名称，你可以查看图可视化，或者直接打印 ModelHistory 对象的输出（两者均已在上文展示）。以下是一些提取特定层信息以及该层实际激活值的示例：

print(model_history['conv2d_3_7'])  # 通过层名提取
# 以下注释掉的行也能提取同一层：
# model_history['conv2d_3'] 可以省略第二个数字（因为严格来说它是冗余的）
# model_history['conv2d_3_7:1'] 冒号表示某层的哪一次传递（这里只有一次）
# model_history['features.6'] 可以通过该层所属的模块来提取
# model_history[7] 整体的第7层

# model_history[-17] 倒数第17层
'''
层 conv2d_3_7，操作 8/24：
	输出张量：形状=(1, 384, 13, 13)，类型=torch.float32，大小=253.5 KB
		张量([[ 0.0503, -0.1089, -0.1210, -0.1034, -0.1254],
        [ 0.0789, -0.0752, -0.0581, -0.0372, -0.0181],
        [ 0.0949, -0.0780, -0.0401, -0.0209, -0.0095],
        [ 0.0929, -0.0353, -0.0220, -0.0324, -0.0295],
        [ 0.1100, -0.0337, -0.0330, -0.0479, -0.0235]])...
	参数：由形状为 (384,) 和 (384, 192, 3, 3) 的参数计算得出；共 663936 个参数（2.5 MB）
	父层：maxpool2d_2_6
	子层：relu_3_8
	函数：conv2d（grad_fn=ConvolutionBackward0）
	在模块中计算：features.6
	耗时：5.670E-04s
	模块输出：features.6
	底层模块输出：features.6
	查找键：-17、7、conv2d_3_7、conv2d_3_7:1、features.6、features.6:1
'''

# 你可以通过 activation 字段提取某一层的实际输出激活值：
print(model_history['conv2d_3_7'].activation)
'''
tensor([[[[-0.0867, -0.0787, -0.0817,  ..., -0.0820, -0.0655, -0.0195],
          [-0.1213, -0.1130, -0.1386,  ..., -0.1331, -0.1118, -0.0520],
          [-0.0959, -0.0973, -0.1078,  ..., -0.1103, -0.1091, -0.0760],
          ...,
          [-0.0906, -0.1146, -0.1308,  ..., -0.1076, -0.1129, -0.0689],
          [-0.1017, -0.1256, -0.1100,  ..., -0.1160, -0.1035, -0.0801],
          [-0.1006, -0.0941, -0.1204,  ..., -0.1146, -0.1065, -0.0631]]...
'''

如果你不想保存所有层的激活值（例如为了节省内存），可以在调用 log_forward_pass 时使用 layers_to_save 参数指定要保存的层；你可以像上面那样通过索引方式指定层，也可以传入一个子字符串来筛选层（例如，'conv' 会提取所有卷积层）：

# 提取 conv2d_3_7、'features' 模块的输出、倒数第五层以及所有线性层（即 fc 层）：
model_history = tl.log_forward_pass(alexnet, x, vis_mode='unrolled',
                                    layers_to_save=['conv2d_3_7', 'features', -5, 'linear'])
print(model_history.layer_labels)
'''
['conv2d_3_7', 'maxpool2d_3_13', 'linear_1_17', 'dropout_2_19', 'linear_2_20', 'linear_3_22']
'''

TorchLens 的主要功能是 log_forward_pass；其余功能包括：

1. get_model_metadata：用于检索所有模型元数据而无需保存任何激活值（例如，用于确定要保存哪些层；请注意，这与调用 log_forward_pass 并设置 layers_to_save=None 相同）
2. show_model_graph：用于可视化模型图结构，同样无需保存任何激活值
3. validate_model_activations：运行一个流程来检查激活值是否正确：具体来说，它会先执行一次前向传播并保存所有中间激活值，然后从每个中间层重新开始前向传播，并验证结果是否与真实输出一致。此外，它还会检查用随机噪声替换保存的激活值是否会生成错误的输出。如果此函数返回 False（即保存的激活值有误），请通过电子邮件 (johnmarkedwardtaylor@gmail.com) 或本 GitHub 页面告知我问题详情，我将更新 TorchLens 以修复该问题。

以上就是全部内容。TorchLens 目前仍在积极开发中，目标是能够毫无例外地兼容任何 PyTorch 模型。截至撰写本文时，它已在超过 700 种图像、视频、音频、多模态和语言模型上进行了测试，涵盖前馈网络、循环神经网络、Transformer 以及图神经网络等。

其他功能

• 你可以使用 log_forward_pass 的 vis_nesting_depth 参数以不同嵌套深度可视化模型；例如，这里展示了 GoogLeNet 的“inception”模块在不同嵌套深度下的视图：

• 一项实验性功能是可以不仅提取模型所有操作的激活值，还可以提取反向传播中的梯度（这些梯度可以基于任意中间层计算，而不仅仅是模型的输出），并可视化反向传播的路径（如下方蓝色箭头所示）。有关如何实现此功能的说明，请参阅 CoLab 教程。

• 你可以通过 func_call_stack 字段查看运行模型所用的确切代码。每一项都是一个 FuncCallLocation 对象，包含清晰的表示和源代码上下文：

print(model_history['conv2d_3'].func_call_stack[0])
'''
FuncCallLocation：
  文件：/usr/local/lib/python3.10/dist-packages/torchvision/models/alexnet.py
  行：48
  函数：forward
  代码：
          x = self.features(x)
          x = self.avgpool(x)
    --->  x = self.classifier(x)
          return x
'''

计划中的功能

1. 在未来，我计划增加一项功能，不仅可以保存激活值，还可以对它们进行反事实干预（例如，如果这些参数不同或使用了不同的非线性激活函数，输出会如何变化）。请告诉我你是否觉得这项功能有用，以及你希望具体实现什么样的功能。
2. 我计划添加一个选项，允许只可视化模型的一个子模块，而不是一次性展示整个图结构。

其他值得了解的工具包

我们的目标是让 TorchLens 完全解决从深度神经网络中提取激活值和元数据、并可视化其结构的问题，从而让任何人都不再需要为此操心。然而，它有意省略了一些功能：例如，没有用于加载模型或刺激数据、也没有用于分析提取出的激活值的函数。这部分原因在于，我们无法预知你可能想对这些激活值做哪些操作，也无法预见所有你可能想要研究的模型类型；另一方面，目前已经有一些非常优秀的工具包可以完成这些任务。以下是一些推荐的工具包——如果你觉得我遗漏了什么，请告诉我！

• Cerbrec：一款用于交互式可视化和调试深度神经网络的程序（底层使用了 TorchLens 来提取 PyTorch 模型的计算图！）
• ThingsVision：提供了强大的视觉模型加载、刺激数据加载以及提取激活值后分析的功能。
• Net2Brain：与 ThingsVision 类似，同样具备出色的端到端功能，并且还支持将提取的激活值与神经科学数据进行比较。
• surgeon-pytorch：易于使用的模型激活值提取功能，同时还支持基于中间层激活值定义损失函数来训练模型。
• deepdive：在加载和基准测试多种不同模型方面表现出色。
• torchvision feature_extraction 模块：可以从具有静态计算图的模型中提取激活值。
• rsatoolbox3：一个完整的解决方案，可用于对深度神经网络的激活值和脑数据进行表征相似性分析。

致谢

TorchLens 的开发得益于与 Nikolaus Kriegeskorte、George Alvarez、Alfredo Canziani、Tal Golan 以及哥伦比亚大学视觉推理实验室的深入讨论。同时，也感谢 Kale Kundert 提供的有益讨论及代码贡献，使得 TorchLens 能够兼容 PyTorch Lightning。所有网络可视化均使用 Graphviz 生成。Logo 由 Nikolaus Kriegeskorte 设计。

引用 TorchLens

若需引用 TorchLens，您可以参考介绍该工具包的论文，并在发现 TorchLens 有用时为本仓库点亮一颗星：

Taylor, J., Kriegeskorte, N. 使用 TorchLens 提取并可视化 PyTorch 模型的隐藏激活值和计算图。Sci Rep 13, 14375 (2023). https://doi.org/10.1038/s41598-023-40807-0

联系方式

由于 TorchLens 仍处于积极开发阶段，我们非常期待您的反馈。如有任何问题、意见或建议，或者您有兴趣开展合作，请通过以下方式与我们联系：

• 邮箱：johnmarkedwardtaylor@gmail.com
• Twitter：twitter.com/johnmark_taylor
• GitHub 仓库的 Issues 或 Discussions 页面。

TorchLens 快速上手指南

TorchLens 是一个强大的 PyTorch 工具，旨在无需修改模型代码的情况下，一键提取模型中每一个中间操作的激活值（Activations），并提供直观的计算图可视化及详细的元数据。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• Python 版本：兼容主流 Python 3.x 版本。
• PyTorch 版本：需要 1.8.0 或更高版本。
• 系统依赖：必须安装 Graphviz 用于生成网络结构可视化图像。
  • Ubuntu/Debian: sudo apt install graphviz
  • macOS (Homebrew): brew install graphviz
  • Windows: 请下载并安装 Graphviz 安装包，并将 bin 目录添加到系统环境变量 PATH 中。

安装步骤

推荐使用国内镜像源加速安装过程。

1. 安装系统依赖 Graphviz（如已安装可跳过）：

   sudo apt install graphviz
   

2. 通过 pip 安装 TorchLens（使用清华源）：

   pip install torchlens -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

TorchLens 的核心功能是 log_forward_pass。只需一行代码，即可记录模型的前向传播过程，获取所有中间层的激活值和计算图。

最简单的使用示例

以下示例展示了如何对预训练的 AlexNet 模型进行完整的前向传播日志记录：

import torch
import torchvision
import torchlens as tl

# 1. 加载模型和输入数据
alexnet = torchvision.models.alexnet()
x = torch.rand(1, 3, 224, 224)

# 2. 执行前向传播记录
# layers_to_save='all': 保存所有层的激活值
# vis_mode='unrolled': 展开显示计算图
model_history = tl.log_forward_pass(alexnet, x, layers_to_save='all', vis_mode='unrolled')

# 3. 查看摘要信息
print(model_history)

# 4. 提取特定层的激活值
# 可以通过层名称（如 'conv2d_3_7'）、模块名或索引位置访问
layer_info = model_history['conv2d_3_7']
activations = layer_info.activation

print(f"激活值形状：{activations.shape}")

关键功能说明

• 自动可视化：运行上述代码后，TorchLens 会自动生成并显示模型的计算结构图（需支持图形显示的终端或 Notebook 环境）。
• 灵活索引：您可以通过多种方式访问特定层的数据：
  • 按层名称：model_history['conv2d_3_7']
  • 按模块输出：model_history['features.6']
  • 按顺序索引：model_history[7] （第 7 层）或 model_history[-1] （最后一层）
• 节省内存：如果不需要保存所有层，可通过 layers_to_save 参数指定特定层或关键词（如 ['conv', 'linear']）来过滤。

# 仅保存卷积层和名为 'features' 的模块输出
model_history = tl.log_forward_pass(alexnet, x, layers_to_save=['conv', 'features'])