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参数信息为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可视化一个简单的神经网络