pytorch中模型结构图的可视化

1. 使用summary输出模型结构

• 安装torchsummary：pip install torchsummary
• 使用：
  from torchsummary import summary
summary(your_model, input_size=(channels, H, W), device='cpu')
• 参考文章：使用torchsummary打印torch每层形状：https://www.cnblogs.com/yyf2019/p/11674266.html

2. 使用 torchviz 输出模型结构图

• 安装Graphviz：sudo apt-get install graphviz
• 在虚拟环境下安装： pip install graphviz和 pip install torchviz
• 使用

from torchviz import make_dot

# ...... 

x = Variable(torch.randn(1, 3, 32, 32))
vis_graph = make_dot(model(x)) # 或者 vis_graph = make_dot(model(x), params=dict(model.named_parameters()))
vis_graph.view('VGGNet.gv')  # 保存的文件名，会自动生成 xxx.gv 和 xxx.pdf

• 参考文章：https://blog.csdn.net/gulaixiangjuejue/article/details/108631734

3. 其他的可视化方法：

• 使用微软的tensorwatch （只能在jupyter notebook中使用，个人最喜欢这种方式）

• 使用netron可视化工具（.pt 或者是 .pth 文件）

• 参考博客：

pytorch 网络结构可视化方法汇总（三种实现方法详解）