0. 简介

PyTorch是一个深度学习框架,它使用张量(tensor)作为核心数据结构。在可视化PyTorch模型时,了解每个张量运算的意义非常重要。张量运算作为神经网络模型中的基本操作。它们用于处理输入数据、执行权重更新和生成预测结果。同时张量运算还用于计算损失函数。损失函数衡量了模型预测与真实标签之间的差异。通过使用张量运算,可以计算出模型的预测结果与真实标签之间的差异,并将其最小化。所以一款能够可视化任何PyTorch模型的张量显示开源项目非常重要。这里是该项目的Github地址。

1. 了解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_opt='rolled')
print(model_history['linear_1_1:2'].tensor_contents)  # second pass of first linear layer

'''
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]])
'''


这是一个非常复杂的变压器模型(swin_v2_b),其前向传递过程中涉及了1932个操作;我们也可以获取每个操作的保存输出。

2. 安装TorchLens

要安装TorchLens,请先安装graphviz(用于生成网络可视化),然后使用pip安装TorchLens:

sudo apt install graphviz
pip install torchlens

TorchLens与PyTorch的1.8.0及以上版本兼容

3. 如何使用TorchLens

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_opt='unrolled')
print(model_history)

'''
Log of AlexNet forward pass:
    Model structure: purely feedforward, without branching; 23 total modules.
    24 tensors (4.8 MB) computed in forward pass; 24 tensors (4.8 MB) saved.
    16 parameter operations (61100840 params total; 248.7 MB).
    Random seed: 3210097511
    Time elapsed: 0.288s
    Module Hierarchy:
        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
    Layers:
        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对象中提取有关给定层的信息,包括其激活和有用的元数据:

  • 通过层的名称(按照“conv2d_3_7”表示第3个卷积层,总共是第7个层)
  • 通过该层作为输出的模块的名称(例如,“features”或“classifier.3”)
  • 通过层的序号位置(例如,第2个层为2,倒数第五个为-5;这里输入和输出也被视为层)
    要快速确定这些名称,您可以查看图形可视化或打印ModelHistory对象的输出(如上所示)。以下是一些示例,说明如何提取有关特定层的信息,以及如何提取该层的实际激活值:
print(model_history['conv2d_3_7'])  # pulling out layer by its name 
# The following commented lines pull out the same layer:
# model_history['conv2d_3'] you can omit the second number (since strictly speaking it's redundant)
# model_history['conv2d_3_7:1'] colon indicates the pass of a layer (here just one)
# model_history['features.6'] can grab a layer by the module for which it is an output
# model_history[7] the 7th layer overall
# model_history[-17] the 17th-to-last layer
'''
Layer conv2d_3_7, operation 8/24:
    Output tensor: shape=(1, 384, 13, 13), dype=torch.float32, size=253.5 KB
        tensor([[ 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]])...
    Params: Computed from params with shape (384,), (384, 192, 3, 3); 663936 params total (2.5 MB)
    Parent Layers: maxpool2d_2_6
    Child Layers: relu_3_8
    Function: conv2d (gradfunc=ConvolutionBackward0) 
    Computed inside module: features.6
    Time elapsed:  5.670E-04s
    Output of modules: features.6
    Output of bottom-level module: features.6
    Lookup keys: -17, 7, conv2d_3_7, conv2d_3_7:1, features.6, features.6:1
'''

# You can pull out the actual output activations from a layer with the tensor_contents field: 
print(model_history['conv2d_3_7'].tensor_contents)
'''
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’将提取所有的卷积层)

# Pull out conv2d_3_7, the output of the 'features' module, the fifth-to-last layer, and all linear (i.e., fc) layers:
model_history = tl.log_forward_pass(alexnet, x, vis_opt='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']
'''

4. 其他功能

我们可以使用vis_nesting_depth参数来可视化不同层次的嵌套深度的模型,例如,在这里您可以看到GoogLeNet的一个“inception”模块在不同层次的嵌套深度下的情况:

此外这个项目还可以提取反向传播的梯度(可以基于任何中间层计算,而不仅仅是模型的输出),并且可以可视化反向传播的路径(如下图所示的蓝色箭头)。请参考CoLab教程了解如何实现这一功能。