ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

ubuntu20下使用 torchviz可视化计算图

2026/9/19 15:13:58 拓冰建站 浏览量
ubuntu20下使用 torchviz可视化计算图
  1. 安装 torchviz:
pip install torchviz
  1. 示例代码:下面是一个简单的示例代码,展示如何使用 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. 生成计算图:运行上述代码后,将会在当前目录下生成一个名为 computational_graph.png的图片文件,其中包含了模型的计算图结构。

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

  1. 终端输入 sudo apt-get install graphviz
  2. 添加环境变量,which dot #查看安装路径,export PATH=/which dot 结果:$PATH
    在这里插入图片描述