PyTorch 网络可视化实战：Jupyter Notebook 与 TensorWatch 的完美结合

1. 为什么我们需要在Jupyter里“看见”神经网络？

搞深度学习，尤其是用PyTorch，我们经常把网络结构写在一行行代码里。nn.Conv2d、nn.Linear、nn.ReLU……这些层堆叠起来，就成了一个复杂的模型。但说实话，光看代码，你脑子里能立刻浮现出这个网络长什么样吗？有多少层？数据是怎么流动的？层与层之间是怎么连接的？对于新手，甚至是有经验的开发者，单凭想象去理解一个复杂网络，都是一件挺费劲的事儿。

这就好比给你一堆乐高积木的零件清单，让你在脑子里拼出完整的城堡。清单再详细，也不如直接给你一张搭建图纸来得直观。网络可视化，就是这张“搭建图纸”。它能把你用代码定义的抽象网络，变成一个清晰、直观的图形，让你一眼就能看清模型的“骨架”。

那么，为什么特别强调在 Jupyter Notebook 里做这件事呢？我自己的体会是，Jupyter 是一个探索和实验的绝佳环境。你写一段代码，立刻就能看到结果，这种即时反馈对学习和调试来说太重要了。想象一下，你刚定义了一个新的残差块，或者修改了某个通道数，马上就能在同一个Notebook里看到图形化的结构变化，这种体验是命令行或者纯脚本文件难以比拟的。它把代码编写、模型构建和结果可视化无缝地衔接在了一起，形成了一个高效的工作流闭环。

而 TensorWatch 这个工具，就是专门为这种交互式探索而生的。它最初由微软研究院开发，设计理念就是让研究者能方便地监控和可视化PyTorch模型的训练过程，其中自然也包括了模型结构本身。相比于其他一些可视化方案，TensorWatch与Jupyter的集成度更高，调用起来更像是一个“魔法命令”，几行代码就能出图，非常符合我们快速验证、即时观察的需求。

当然，光有TensorWatch还不够，它需要一个强大的“画图引擎”来把计算图渲染成漂亮的图形，这就是 Graphviz。Graphviz是一个老牌的开源图形可视化软件，擅长处理点线图关系。TensorWatch负责解析PyTorch模型，生成Graphviz能理解的描述语言，然后由Graphviz渲染成最终的SVG或图片。所以，我们这个“完美结合”的实质，是 Jupyter提供舞台，TensorWatch担任导演和编剧，Graphviz则是那位技艺高超的美术师，三者缺一不可。

2. 从零开始：搭建你的可视化工作台

好了，道理讲明白了，咱们直接上手。我把自己配置环境时踩过的坑和总结的最佳实践都揉在这里，你跟着一步步来，大概率能一路绿灯。

2.1 核心环境准备：版本匹配是头等大事

首先我得强调一个血泪教训：版本兼容性。PyTorch、TensorWatch、Graphviz，甚至Python本身，版本之间都有一些“小脾气”。如果随便装，很可能遇到各种稀奇古怪的报错。我推荐一套经过我多次实测、非常稳定的组合，这也是原文中提到的“黄金配置”：

• Python: 3.6.5 或 3.7。不建议直接用最新的3.9+，有些老版本库的兼容性支持可能还没跟上。
• PyTorch: 1.2.0。这是一个非常经典的稳定版本。
• torchvision: 0.4.0，与PyTorch 1.2.0配套。
• tensorwatch: 0.8.7。注意，这个库后续的维护似乎不太活跃，0.8.7版本与上述环境配合良好。
• Graphviz (软件): 从官网下载的稳定版即可。
• graphviz (Python包): 通过pip安装。
• pydot: 1.4.2。这是一个关键的连接器，版本不对容易出问题。

怎么管理这些版本呢？我强烈建议使用 Anaconda 来创建独立的虚拟环境。打开你的Anaconda Prompt（或者终端，如果你用miniconda），执行下面的命令来创建一个新环境：

conda create -n pytorch_viz python=3.6.5
conda activate pytorch_viz

然后在这个环境里安装PyTorch。由于是历史版本，我们最好去PyTorch官网的历史版本库找安装命令，或者直接用pip指定版本安装：

pip install torch==1.2.0 torchvision==0.4.0 -f https://download.pytorch.org/whl/torch_stable.html

这样就奠定了最坚实的基础。用虚拟环境的好处是，这个配置好的环境就像是一个干净的“实验室”，不会和你电脑上其他项目的环境冲突，玩坏了删掉重来也很方便。

2.2 搞定Jupyter Notebook的内核切换

装好PyTorch环境后，你直接打开桌面的Jupyter Notebook，会发现新建文件时，可用的内核可能只有基础的“Python 3”。这可不是我们刚配好的pytorch_viz环境。我们需要把这个环境“注册”到Jupyter里。

操作很简单，但步骤要对：

1. 确保你已经激活了pytorch_viz环境（命令行前面有(pytorch_viz)字样）。
2. 安装ipykernel，这是Jupyter内核的依赖包：pip install ipykernel。
3. 将当前环境添加到Jupyter的可选内核列表中：

   python -m ipykernel install --user --name pytorch_viz --display-name "PyTorch 1.2 (Visual)"
   

   这里的--name后面跟的是你conda环境的实际名称（pytorch_viz），--display-name是你在Jupyter界面上看到的漂亮名字，可以自己起个好认的。

完成之后，关闭所有已经打开的Jupyter Notebook页面和服务器。重新启动Jupyter，点击“New”，你就会惊喜地发现多了一个“PyTorch 1.2 (Visual)”的选项。选择它，你在这个Notebook里使用的就是咱们配好的、包含了特定版本PyTorch的完整环境了。这一步是很多新手会忽略的，结果就是一直在基础环境里折腾，怎么也导不入torch包。

2.3 安装可视化“三剑客”：Graphviz软件与Python包

接下来安装我们的“美术师”Graphviz。这里有个关键点：你需要安装两个东西，一个是Graphviz软件本身（用来画图的引擎），另一个是Python的graphviz包（用来调用这个引擎的接口）。

首先安装Python包（在pytorch_viz环境下）：

pip install graphviz

这个通常很快。

然后安装Graphviz软件。你需要去Graphviz的官网下载安装程序。对于Windows用户，原文中给出的下载链接可能已经变化，更稳妥的方法是访问Graphviz的官方下载页（https://graphviz.org/download/），选择适合你操作系统的版本。下载完成后，运行安装程序。安装过程中，请务必留意一个选项：“Add Graphviz to the system PATH for all users”或类似表述，一定要勾选上！ 这步至关重要，它让系统知道Graphviz的可执行文件（比如dot命令）在哪里，否则Python包找不到它，后面就会报错。

安装完成后，你可以验证一下。打开一个新的命令行（需要重启命令行窗口以使环境变量生效），输入：

dot -V

如果显示了Graphviz的版本信息，比如dot - graphviz version 2.50.0，那就说明软件安装和环境变量配置都成功了。

2.4 安装TensorWatch及其他辅助依赖

“美术师”就位了，现在请“导演”TensorWatch登场。回到激活的pytorch_viz环境，我们一次性安装所有必需的Python包：

pip install tensorwatch==0.8.7 pydot==1.4.2 scikit-learn==0.24.2 pandas==1.1.5

这里解释一下每个包的作用：

• tensorwatch==0.8.7: 主角，负责模型可视化。
• pydot==1.4.2: 另一个处理Graphviz文件的Python接口，TensorWatch内部会用到它。版本锁定为1.4.2是为了避免兼容性问题。
• scikit-learn和pandas: 虽然不是可视化直接必需的，但TensorWatch的一些高级功能或示例可能会用到它们，装上以备不时之需，版本按原文推荐即可。

现在，你的可视化“全家桶”就已经在虚拟环境里准备完毕了。

3. 实战演练：让你的第一个网络“跃然纸上”

环境配好了，手有点痒了吧？咱们马上在Jupyter Notebook里跑一个真实的例子。打开Jupyter，选择我们刚创建的“PyTorch 1.2 (Visual)”内核，新建一个Notebook。

3.1 第一个可视化测试与经典报错解决

我们先用一个经典的、现成的模型来测试。在第一个单元格里，输入以下代码：

from torchvision.models import vgg16  # 导入经典的VGG16模型
from tensorwatch import draw_model    # 导入TensorWatch的绘图函数

# 实例化一个VGG16模型
mynet = vgg16()
# 绘制模型结构。第二个参数是输入张量的形状 [batch_size, channels, height, width]
draw_model(mynet, [1, 3, 224, 224])

满怀期待地按下Shift+Enter运行，结果很可能不是一张漂亮的图，而是一段红色的报错信息，核心内容可能是：'Dot' object has no attribute '_repr_svg_'。

别慌，这个几乎是每个第一次用TensorWatch的人都会遇到的“经典关卡”。问题出在pydot库与Python新版本的兼容性上。解决它需要一点“小手术”。

解决办法如下：

1. 找到你当前Python环境（pytorch_viz）的site-packages目录。一个简单的查找方法是，在Jupyter Notebook里新建一个单元格，运行：

   import site; site.getsitepackages()
   

   它会输出路径列表，通常第一个就是。比如可能是C:\Users\你的用户名\Anaconda3\envs\pytorch_viz\Lib\site-packages 或者 /home/username/anaconda3/envs/pytorch_viz/lib/python3.6/site-packages。
2. 在这个site-packages目录下，找到名为tensorwatch的文件夹，进入，再找到pytorch_draw_model.py这个文件。
3. 用任何文本编辑器（比如VS Code、Notepad++，甚至记事本）打开这个文件。
4. 找到第13行左右（不同版本可能略有差异），你会看到一行代码类似：

   return self.dot._repr_svg_() if hasattr(self.dot, ‘_repr_svg_’) else self.dot
   

   或者直接就是 return self.dot._repr_svg_()。
5. 将这一行修改为：

   return self.dot.create_svg().decode()
   

   这个修改的目的是用pydot库更通用的create_svg()方法来生成SVG格式数据，并用decode()将其转换为字符串。
6. 保存文件。

修改完成后，有一个极其重要的步骤：回到Jupyter Notebook，重启内核！ 点击菜单栏的“Kernel” -> “Restart”。因为Python已经加载了旧的模块代码到内存中，不重启内核，你的修改是不会生效的。

内核重启后，再次运行刚才的绘图代码。这次，你应该能看到一个非常长的、纵向展开的网络结构图了！图中，每个方框代表一个计算层，箭头代表数据流向，层名、输入输出形状等信息都标注得清清楚楚。你可以拖动滚动条，从输入层一直浏览到输出层，对整个VGG16的层次堆积有一个前所未有的直观认识。

3.2 可视化你自己的自定义网络

用现成模型测试成功，只是第一步。我们更想看的，是自己设计的网络。下面我们来定义一个简单的卷积神经网络（CNN），并可视化它。

import torch
import torch.nn as nn
from tensorwatch import draw_model

# 定义一个简单的CNN
class MySimpleCNN(nn.Module):
    def __init__(self):
        super(MySimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.flatten = nn.Flatten()
        # 假设经过两次池化后，特征图大小是 56x56 (如果输入是224x224)
        self.fc1 = nn.Linear(in_features=32 * 56 * 56, out_features=128)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(in_features=128, out_features=10) # 10个分类

    def forward(self, x):
        x = self.pool1(self.relu1(self.conv1(x)))
        x = self.pool2(self.relu2(self.conv2(x)))
        x = self.flatten(x)
        x = self.relu3(self.fc1(x))
        x = self.fc2(x)
        return x

# 实例化并可视化
my_cnn = MySimpleCNN()
# 注意输入形状要和网络第一层匹配
draw_model(my_cnn, [1, 3, 224, 224])

运行这段代码，你就能看到MySimpleCNN这个自定义类的结构图了。通过这个图，你可以清晰地验证：

• 卷积层conv1的输入输出通道数是否正确。
• 池化层pool1之后，特征图的尺寸是否如你预期那样减半。
• Flatten层是如何将多维特征图“压平”成一维向量的。
• 全连接层fc1的输入特征数是否等于Flatten后的向量长度（这里是32*56*56）。这里非常容易出错，如果计算错了，图可能画不出来，或者模型运行时才会报维度错误。可视化帮你提前发现了这个问题。

3.3 探索TensorWatch的更多可视化技巧

基础的绘图功能掌握了，TensorWatch还有一些小技巧可以让可视化更高效。

控制图的复杂度：对于非常深的网络（如ResNet-152），默认绘制的图会巨大无比，在Jupyter里查看可能很卡。draw_model函数有一个depth参数可以控制展开的深度。例如draw_model(model, [1,3,224,224], depth=2)可能只显示最高几层的模块结构，适合看宏观架构。

结合模型摘要：可视化是看“形”，有时候我们还想看“数”，比如每一层的参数数量。可以结合torchsummary库（需要单独安装：pip install torchsummary）来使用：

from torchsummary import summary

summary(my_cnn, input_size=(3, 224, 224), device='cpu')

这段代码会在下方输出一个表格，详细列出每一层的类型、输出形状、参数数量。左边看图理解结构，右边看表掌握参数量，两者结合，你对模型的理解就立体起来了。

可视化子模块：如果你的网络是由几个大的子模块（如Backbone、Neck、Head）组成的，你可以单独可视化其中一个部分，让图形更聚焦：

# 假设my_cnn有一个特征提取的子模块叫`features`
draw_model(my_cnn.features, [1, 3, 224, 224])

4. 避坑指南与最佳实践

走通了整个流程，我再来分享几个我踩过坑后总结的经验，能让你用得更顺手。

关于环境变量：Graphviz安装后如果dot命令还是找不到，除了检查安装时是否勾选添加PATH，还可以手动添加。将Graphviz安装目录下的bin文件夹路径（例如C:\Program Files\Graphviz\bin）添加到系统的环境变量PATH中，然后重启命令行和Jupyter Notebook服务器。

关于内核重启：在Jupyter中，每当你安装了新的Python包，或者修改了像pytorch_draw_model.py这样的库文件，都必须重启内核，新的更改才会被加载。这是一个非常重要的习惯。

版本问题的变通：如果你因为项目原因必须使用更新的PyTorch版本（比如1.7+），而TensorWatch 0.8.7可能兼容性不佳。除了尝试寻找更晚的TensorWatch版本外，也可以考虑PyTorch官方自带的torchviz库，或者另一个非常强大的工具 Netron。Netron支持离线打开模型文件（如PyTorch的.pt或.pth文件），提供交互式查看，功能非常专业。你可以用torch.jit.trace或torch.jit.script导出模型，然后用Netron打开查看。

图形渲染太慢或太大：如果网络极其复杂，生成SVG图可能会让浏览器卡顿。你可以考虑将可视化输出到文件，然后用图片查看器打开。虽然TensorWatch的draw_model主要针对Jupyter，但你可以通过一些方法捕获其输出的SVG字符串，保存为文件。或者，直接使用graphviz库的原始接口来生成PDF或PNG，虽然代码会稍微复杂一点，但对于超大模型更稳定。

用于教学与沟通：最后，别忘了这个工具的巨大潜力——用于教学和团队沟通。在文档、技术报告甚至PPT里，插入一张自动生成的、清晰准确的网络结构图，比用绘图软件手动画一个要专业和高效得多，而且能保证与代码完全一致。我经常在项目文档的开头，先放一张核心模型的可视化图，让读者（包括一段时间后的我自己）能快速建立起对项目的整体认知。