pytorch 模型可视化

pytorch的模型结构可视化方法:

(1)使用tensorboardX(不太直观)

(2)使用graphviz加上torchviz (依赖于graphviz和GitHub第三方库torchviz)

(3)使用微软的tensorwatch (只能在jupyter notebook中使用,个人最喜欢这种方式)

(4)使用netron可视化工具(.pt 或者是 .pth 文件)

一、使用tensorboardX

使用tensorboardX必须要安装tensorboard才行,可能会存在一些版本的匹配问题,下面的版本是亲测有效的。

pytorch0.4.1+tensorboard1.7.0+tensorboardX1.4

pytorch1.0.1+tensorboard1.14.0+tensorboardX1.8

上面这两个都是可行的。

本文以AlexNet为例,鉴于torchvision已经实现了AlexNet模型,就不再自己编写,直接导入即可。

import torch
from torchvision.models import AlexNet
 
from tensorboardX import SummaryWriter
 
 
x=torch.rand(8,3,256,512)
model=AlexNet()
 
with SummaryWriter(comment='AlexNet') as w:
    w.add_graph(model, x)  # 这其实和tensorflow里面的summarywriter是一样的。

上面的代码运行结束后,会在当前目录生成一个叫runs的文件夹,runs文件夹里面会有一个文件夹Jul22_18-03-19_WH-PC19012AlexNet,里面存储了可视化所需要的日志信息。用cmd进入到runs文件夹所在的目录中(路劲中不能有中文),然后cmd中输入:

tensorboard --logdir Jul22_18-03-19_WH-PC19012AlexNet

二、使用graphviz+torchviz来可视化模型

首先安装这两个依赖包

pip install graphviz  # 安装graphviz
pip install git+https://github.com/szagoruyko/pytorchviz  # 通过git安装torchviz

第一步:加载并运行一个模型

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)
# g=make_dot(y, params=dict(model.named_parameters()))
#g = make_dot(y, params=dict(list(model.named_parameters()) + [('x', x)]))

第三步:保存模型,以PDF格式保存

# 这两种方法都可以
# g.view() # 会生成一个 Digraph.gv.pdf 的PDF文件
g.render('espnet_model', view=False) # 会自动保存为一个 espnet.pdf,第二个参数为True,则会自动打开该PDF文件,为False则不打开

另外,我还可以查询整个模型的参数量信息,代码如下

# 查看模型的参数信息
 
params = list(model.parameters())
k = 0
for i in params:
        l = 1
        print("该层的结构:" + str(list(i.size())))
        for j in i.size():
                l *= j
        print("该层参数和:" + str(l))
        k = k + l
print("总参数数量和:" + str(k))
'''
该层的结构:[64, 3, 11, 11]
该层参数和:23232
该层的结构:[64]
该层参数和:64
该层的结构:[192, 64, 5, 5]
该层参数和:307200
该层的结构:[192]
该层参数和:192
该层的结构:[384, 192, 3, 3]
该层参数和:663552
该层的结构:[384]
该层参数和:384
该层的结构:[256, 384, 3, 3]
该层参数和:884736
该层的结构:[256]
该层参数和:256
该层的结构:[256, 256, 3, 3]
该层参数和:589824
该层的结构:[256]
该层参数和:256
该层的结构:[4096, 9216]
该层参数和:37748736
该层的结构:[4096]
该层参数和:4096
该层的结构:[4096, 4096]
该层参数和:16777216
该层的结构:[4096]
该层参数和:4096
该层的结构:[1000, 4096]
该层参数和:4096000
该层的结构:[1000]
该层参数和:1000
总参数数量和:61100840
'''

三、通过tensorwatch+jupyter notebook来实现

import torch
import tensorwatch as tw
from lanenet_model.blocks import ESPNet_Encoder # 这是我自己定义的一个网络
 
# 其实就两句话
model=ESPNet_Encoder()
tw.draw_model(model, [1, 3, 512, 256])

四 打印网络结构

import torch
import torch.nn as nn
from torchsummary import summary
from tensorboardX import SummaryWriter

# unetmodel = Model().to(device)
# summary(unetmodel,(3,128,128),batch_size=2)

# ------------------------------------------
# unet = Model()
# # print(unet)
# for i in unet._modules.items():
#     print(i)
# ------------------------------------------

# x=torch.rand(2,3,128,128).to(device)
# model = Model().to(device)
# with SummaryWriter(comment='Unet') as w:
#     w.add_graph(model,x)

你可能感兴趣的:(pytorch,深度学习,python)