0、引言

深度学习中,模型结构一般都很复杂,是否有一款好用的工具可以讲模型结构可视化呢?本文介绍一款支持pytorch框架的网络模型可视化工具torchvis。可以将网络结构输出到一个pdf文件中。

1、安装过程

1.1 首先需要下载graphvis。下载链接,选择自己的系统所对应的版本下载即可。
如下图所示,我下载的是最新的。
在这里插入图片描述
1.2 环境变量设置

在这里插入图片描述
双击path,进入新建如下图所示。
在这里插入图片描述

将graphvis的安装路径 加入环境变量。(不知道怎么设置环境变量,可以自行百度)

  1. 检验是否安装成功

命令行输入:dot -version

2、python安装graphvis

pip install  graphvis

3、代码测试

import torch
from torchvision.models import AlexNet
from torchviz import make_dot
 
x=torch.randn(8,3,256,512).requires_grad_(True)
model=AlexNet()
y=model(x)
g = make_dot(y, params=dict(list(model.named_parameters()) + [('x', x)]))
g.render('AlexNet', view=True) 

图片太大,展示部分结果:
在这里插入图片描述

Logo

欢迎加入 MCP 技术社区!与志同道合者携手前行,一同解锁 MCP 技术的无限可能!

更多推荐