ubuntu20下使用 torchviz可视化计算图

1. 安装 torchviz：

pip install torchviz
1

2. 示例代码：下面是一个简单的示例代码，展示如何使用 torchviz 可视化计算图：

python
import torch
from torchviz import make_dot

# 创建一个简单的模型
model = torch.nn.Sequential(
    torch.nn.Linear(2, 2),
    torch.nn.ReLU(),
    torch.nn.Linear(2, 1)
)

# 创建输入数据
x = torch.randn(1, 2)

# 前向传播
y = model(x)

# 可视化计算图
dot = make_dot(y, params=dict(model.named_parameters()))
dot.render("computational_graph", format="png")  # 保存为 PNG 图片
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

3. 生成计算图：运行上述代码后，将会在当前目录下生成一个名为 computational_graph.png的图片文件，其中包含了模型的计算图结构。

出现错误：提示没有dot命令，
解决方法：安装graphviz

1. 终端输入 sudo apt-get install graphviz
2. 添加环境变量，which dot #查看安装路径，export PATH=/which dot 结果:$PATH