方案一:使用netron工具
参考:pytorch模型结构可视化,可显示每层的尺寸 - 知乎 (zhihu.com)
image.png
方案二:tensorwatch+jupyter notebook(限制在jupyter)
效果图:
image.png
方案三:pytorchviz 树形展示
链接:Python库 - Pytorch 模型的网络结构可视化 pytorchviz - AI备忘录 (aiuai.cn)
展示效果:
image.png
代码:
sudo pip install graphviz
# 或
sudo pip install git+https://github.com/szagoruyko/pytorchviz
import torch
from torchvision.models import AlexNet
from torchviz import make_dot
x=torch.rand(8,3,256,512)
model=AlexNet()
y=model(x)
# 调用make_dot()函数构造图对象
g = make_dot(y)
# 保存模型,以PDF格式保存
g.render('Alex_model', view=False)
方案四:torchsummary 文本列表显示
展示效果:
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 64, 64] 1,792
BatchNorm2d-2 [-1, 64, 64, 64] 128
ReLU-3 [-1, 64, 64, 64] 0
MaxPool2d-4 [-1, 64, 32, 32] 0
Conv2d-5 [-1, 64, 32, 32] 36,928
BatchNorm2d-6 [-1, 64, 32, 32] 128
ReLU-7 [-1, 64, 32, 32] 0
MaxPool2d-8 [-1, 64, 16, 16] 0
Conv2d-9 [-1, 64, 16, 16] 36,928
BatchNorm2d-10 [-1, 64, 16, 16] 128
ReLU-11 [-1, 64, 16, 16] 0
MaxPool2d-12 [-1, 64, 8, 8] 0
Conv2d-13 [-1, 64, 8, 8] 36,928
BatchNorm2d-14 [-1, 64, 8, 8] 128
ReLU-15 [-1, 64, 8, 8] 0
MaxPool2d-16 [-1, 64, 4, 4] 0
================================================================
方案五:tensorboard 不建议
展示效果:
image.png