使用TorchLens可视化一个简单的神经网络

  TorchLens：可用于可视化任何PyTorch模型，一个包用于在一行代码中提取和映射PyTorch模型中每个张量运算的结果。TorchLens功能非常强大，如果能够熟练掌握，算是可视化PyTorch模型的一把利剑。本文通过TorchLens可视化一个简单神经网络，算是抛砖引玉吧。

一.定义一个简单神经网络

import torch
import torch.nn as nn
import torch.optim as optim
import torchlens as tl
import os
os.environ["PATH"] += os.pathsep + 'D:/Program Files/Graphviz/bin/'

# 定义神经网络类
class NeuralNetwork(nn.Module): # 继承nn.Module类
    def __init__(self, input_size, hidden_size, output_size):
        super(NeuralNetwork, self).__init__() # 调用父类的构造函数
        # 定义输入层到隐藏层的线性变换
        self.input_to_hidden = nn.Linear(input_size, hidden_size)
        # 定义隐藏层到输出层的线性变换
        self.hidden_to_output = nn.Linear(hidden_size, output_size)
        # 定义激活函数
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 前向传播
        hidden = self.sigmoid(self.input_to_hidden(x))
        output = self.sigmoid(self.hidden_to_output(hidden))
        return output

def NeuralNetwork_train(model):
    # 训练神经网络
    for epoch in range(10000):
        optimizer.zero_grad()  # 清零梯度
        outputs = model(input_data)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播和优化
        optimizer.step()  # 更新参数

        # 每100个epoch打印一次损失
        if (epoch + 1) % 1000 == 0:
            print(f'Epoch [{epoch + 1}/10000], Loss: {loss.item():.4f}')

    return model

def NeuralNetwork_test(model):
    # 在训练后，可以使用模型进行预测
    with torch.no_grad():
        test_input = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.float32)
        predictions = model(test_input)
        predicted_labels = (predictions > 0.5).float()
        print("Predictions:", predicted_labels)

if __name__ == '__main__':
    # 定义神经网络的参数
    input_size = 2  # 输入特征数量
    hidden_size = 4  # 隐藏层神经元数量
    output_size = 1  # 输出层神经元数量

    # 创建神经网络实例
    model = NeuralNetwork(input_size, hidden_size, output_size)

    # 定义损失函数和优化器
    criterion = nn.BCELoss()  # 二分类交叉熵损失
    optimizer = optim.SGD(model.parameters(), lr=0.1)  # 随机梯度下降优化器

    # 准备示例输入数据和标签
    input_data = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.float32)
    labels = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32)

    # model：神经网络模型
    # input_data：输入数据
    # layers_to_save：需要保存的层
    # vis_opt：rolled/unrolled，是否展开循环
    model_history = tl.log_forward_pass(model, input_data, layers_to_save='all', vis_opt='unrolled')  # 可视化神经网络
    print(model_history)
    # print(model_history['input_1'].tensor_contents)
    # print(model_history['input_1'])

    tl.show_model_graph(model, input_data)

    # model = NeuralNetwork_train(model) # 训练神经网络
    # NeuralNetwork_test(model) # 测试神经网络

1.神经网络结构

  输入层包括2个神经元，隐藏层包括4个神经元，输出层包括1个神经元。

2.log_forward_pass

  给定输入x，通过模型运行前向传播，并返回一个包含前向传播日志（层激活和相应的层元数据）的ModelHistory对象。如果vis_opt设置为rolled或unrolled并可视化模型图。

3.show_model_graph

  可视化模型图，而不保存任何激活。

4.查看神经网络模型参数

权重(12)+偏置(5)共计17个参数，如下所示：

二.输出结果分析

1.model_history输出结果

Log of NeuralNetwork forward pass: // 神经网络前向传播日志
	Random seed: 1626722175 // 随机种子
	Time elapsed: 1.742s (1.74s spent logging) // 耗时
	Structure: // 结构
		- purely feedforward, no recurrence // 纯前馈，无循环
		- no branching // 无分支
		- no conditional (if-then) branching // 无条件（if-then）分支
		- 3 total modules // 3个模块
	Tensor info: // 张量信息
		- 6 total tensors (976 B) computed in forward pass. // 前向传播中计算的6个张量（976 B）
		- 6 tensors (976 B) with saved activations. // 6个张量（976 B）保存了激活
	Parameters: 2 parameter operations (17 params total; 548 B) // 参数：2个参数操作（总共17个参数；548 B）
	Module Hierarchy: // 模块层次
		input_to_hidden // 输入到隐藏
		sigmoid:1 // sigmoid:1
		hidden_to_output // 隐藏到输出
		sigmoid:2 // sigmoid:2
	Layers (all have saved activations): // 层（所有层都有保存的激活）
		  (0) input_1        // 输入
		  (1) linear_1_1     // 线性
		  (2) sigmoid_1_2    // sigmoid
		  (3) linear_2_3     // 线性
		  (4) sigmoid_2_4    // sigmoid
		  (5) output_1       // 输出

2.show_model_graph输出结果



（1）总共包含6层

  分别为input_1、linear_1_1、sigmoid_1_2、linear_2_3、sigmoid_2_4和output_1。

（2）总共6个张量

  指的是input_1(160B)、linear_1_1(192B)、sigmoid_1_2(192B)、linear_2_3(144B)、sigmoid_2_4(144B)和output_1(144B)。共计976B。

（3）input_1 4*2(160B)

  4*2表示input_1的shape，而160B指的是该张量在内存中占用空间大小，以字节（B）为单位。知道张量的形状和内存占用情况，对于模型内存管理和优化来说是很有用的信息。其它张量信息如下所示：



（4）共计17参数

  linear_1_1参数信息为42和4，linear_1_1参数信息为14和1，共计17参数，内存占用548B。

三.遇到的问题

1.需要安装和设置graphviz

subprocess.CalledProcessError: Command '[WindowsPath('dot'), '-Kdot', '-Tpdf', '-O', 'graph.gv']' returned non-zero exit status 1.

解决方案是将D:\Program Files\Graphviz\bin添加到系统环境变量PATH中。

2.AlexNet神经网络

因为BP神经网络过于简单，接下来可视化一个稍微复杂点儿的AlexNet神经网络，如下所示：

参考文献：

[1]torchlens_tutorial.ipynb：https://colab.research.google.com/drive/1ORJLGZPifvdsVPFqq1LYT3t5hV560SoW?usp=sharing#scrollTo=W_94PeNdQsUN

[2]Extracting and visualizing hidden activations and computational graphs of PyTorch models with TorchLens：https://www.nature.com/articles/s41598-023-40807-0

[3]torchlens：https://github.com/johnmarktaylor91/torchlens

[4]Torchlens Model Menagerie：https://drive.google.com/drive/folders/1BsM6WPf3eB79-CRNgZejMxjg38rN6VCb

[5]使用TorchLens可视化一个简单的神经网络：github.com/ai408/nlp-engineering/tree/main/20230917_NLP工程化公众号文章/使用torchlens可视化一个简单的神经网络