Pytorch神经网络结构图可视化（Ubuntu+torchviz+graphviz）

Ubuntu环境下

首先安装torch、torchvision、torchviz

pip install torch
pip install torchvision
pip install torchviz

安装graphviz

sudo apt-get install graphviz

如果不安装的话，会有以下报错：

 failed to execute ['dot', '-Tpdf', '-O', 'test'], make sure the Graphviz executables are on your systems' path

大功告成开始测试

import torch
from torchvision import models
from torchviz import make_dot

models = models.resnet18()
x = torch.randn(1,3,224,224)
net_plot = make_dot(models(x),params = dict(models.named_parameters()))
net_plot.view()

生成的图大概是这样的：