- 安装 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
