Pytorch中网络结构的可视化——使用graphviz+torchviz来可视化模型速通

最新推荐文章于 2024-08-14 22:29:49 发布

不易思不逸

最新推荐文章于 2024-08-14 22:29:49 发布

阅读量294

点赞数 9

文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/m0_50877346/article/details/140499891

版权

pip install graphviz torchviz
graphviz官网下载对应的exe安装包 https://graphviz.org/download/ # graphviz-2.50.0 (64-bit) EXE installer
# 测试
import torch
from torchvision.models import AlexNet
from torchviz import make_dot
 
x=torch.rand(8,3,256,512)
model=AlexNet()
y=model(x)
# 这三种方式都可以
g = make_dot(y)
# g=make_dot(y, params=dict(model.named_parameters()))
#g = make_dot(y, params=dict(list(model.named_parameters()) + [('x', x)]))
# 这两种方法都可以
# g.view() # 会生成一个 Digraph.gv.pdf 的PDF文件
g.format = 'png' #设置生成格式， 默认为pdf
g.render('espnet_model', view=False) # 会自动保存为一个 espnet.pdf，第二个参数为True,则会自动打开该PDF文件，为False则不打开