pytorch的模型网络结构的可视化方案

方案一:使用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

你可能感兴趣的:(pytorch的模型网络结构的可视化方案)