TorchLens：可用于可视化任何PyTorch模型，一个包用于在一行代码中提取和映射PyTorch模型中每个张量运算的结果。TorchLens功能非常强大，如果能够熟练掌握，算是可视化PyTorch模型的一把利剑。本文通过TorchLens可视化一个简单神经网络，算是抛砖引玉吧。
一.定义一个简单神经网络
import torchimport torch.nn as nnimport torch.optim as optimimport torchlens as tlimport osos.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 outputdef 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 modeldef 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参数信息为4*2和*4，linear_1_1参数信息为1*4和*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可视化一个简单的神经网络
经验首页