【python】pytorch网络可视化

参考

https://blog.csdn.net/qq_40726937/article/details/106122082

1. graphviz + torchviz

环境安装简单

pip install torchviz
pip install graphviz

代码

import torch
from torchvision import model
from torchviz import make_dot

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

缺点：导出的是pdf