Pytorch的网络结构可视化（tensorboardX）（详细）

版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/xiaoxifei/article/details/82735355
最近刚刚发现一个非常好用的显示模型神器Netron

https://github.com/lutzroeder/Netron

借助这个工具可以像windows的软件一样导入已经训练好的模型加权重即可一键生成

我目前看了下visdom实现pytorch的网络结构查找还是很困难，在stackflow上有很多人使用自己编写的基于matplotlib来实现网络结构可视化适用性也不是很好，后来查找到使用基于tensorboard所开发的tensorboardX可以很方便的实现pytorch网络结构的可视化，因此决定采用这种方式。

1. tensorboardX的简介

tensorboardX的项目路径：https://github.com/lanpa/tensorboardX

tensorboardX是基于tensorboard的思想用来写tensorboard events的工具，可以实现对传统的tensorboard中 scalar，image，figure，histogram，audio，text，graph，onnx_graph等事件进行编写。

tensorboardX同时具有论坛供大家提出问题解决问题   ，论坛地址：https://github.com/lanpa/tensorboardX/wiki

它的支持性是比较好的

2. tensorboardX的使用

tensorboardX的安装以及依赖如下所示：

pip install tensorboard

pip install tensorflow

pip install tensorboardX

tensorboardX的路径下带的有一个规范的demo，可以供大家参考。我这里公布一个我自己测试过的代码，代码来源于：https://blog.csdn.net/sunqiande88/article/details/80155925

import torch
import torch.nn as nn
from tensorboardX import SummaryWriter
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Sequential( #input_size=(1*28*28)
nn.Conv2d(1, 6, 5, 1, 2),
nn.ReLU(), #(6*28*28)
nn.MaxPool2d(kernel_size=2, stride=2), #output_size=(6*14*14)
)
self.conv2 = nn.Sequential(
nn.Conv2d(6, 16, 5),
nn.ReLU(), #(16*10*10)
nn.MaxPool2d(2, 2) #output_size=(16*5*5)
)
self.fc1 = nn.Sequential(
nn.Linear(16 * 5 * 5, 120),
nn.ReLU()
)
self.fc2 = nn.Sequential(
nn.Linear(120, 84),
nn.ReLU()
)
self.fc3 = nn.Linear(84, 10)

# 定义前向传播过程，输入为x
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
# nn.Linear()的输入输出都是维度为一的值，所以要把多维度的tensor展平成一维
x = x.view(x.size()[0], -1)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x

dummy_input = torch.rand(13, 1, 28, 28) #假设输入13张1*28*28的图片
model = LeNet()
with SummaryWriter(comment='LeNet') as w:
w.add_graph(model, (dummy_input, ))
运行该代码后会自动生成一个runs文件夹，并且在文件夹下会有一个对应的event，如下图所示：

此时需要在terminal或者cmd下运行tensorboard --logdir = path

此处千万要注意，如果按照上面的参考文档的方式是会报错的：No graph definition files were found 或者 No definition files were found，总之无法正常显示网络结构图。

此处的path 是event对应的确切，完整的路径

在运行后会出现一个http url，此时需要将该url 拷贝到chrome下即可看到如下所示框图

结构框图如下所示：

此时双击红圈所示的LeNet模块即可看到LeNet的细节信息，如下所示：

至此即可完成使用tensorboardX 对pytorch网络结构的可视化
————————————————
版权声明：本文为CSDN博主「xiaoxifei」的原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/xiaoxifei/article/details/82735355