- 安装 torchviz:
pip install torchviz
- 示例代码:下面是一个简单的示例代码,展示如何使用 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 图片
- 生成计算图:运行上述代码后,将会在当前目录下生成一个名为 computational_graph.png的图片文件,其中包含了模型的计算图结构。
出现错误:提示没有dot命令,
解决方法:安装graphviz
- 终端输入 sudo apt-get install graphviz
- 添加环境变量,which dot #查看安装路径,export PATH=/which dot 结果:$PATH