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

我目前看了下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网络结构的可视化

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

猜你喜欢