pytorch使用tensorboardX进行网络可视化

我们知道，对于pytorch上的搭建动态图的代码的可读性非常高，实际上对于一些比较简单的网络，比如alexnet，vgg阅读起来就能够脑补它们的网络结构，但是对于比较复杂的网络，如unet，直接从代码脑补网络结构可能就比较吃力

tensorflow上的tensorboard进行计算图的可视化可谓是非常成熟了，那么有没有可以可视化pytorch动态图的工具呢？

实际上是有的，前两天介绍了tensorboardX，pytorch上的一个功能强大的可视化工具，他可以直接可视化网络结构

关于如何使用tensorboard，这里还是用一个resnet18来举例子

先贴上代码

#-*-coding:utf-8-*-
import torch
import torchvision
from torch.autograd import Variable
from tensorboardX import SummaryWriter

# 模拟输入数据
input_data = Variable(torch.rand(16, 3, 224, 224))

# 从torchvision中导入已有模型
net = torchvision.models.resnet18()

# 声明writer对象，保存的文件夹，异己名称
writer = SummaryWriter(log_dir='./log', comment='resnet18')
with writer:
    writer.add_graph(net, (input_data,))

torchvision工具包是pytorch自带的强大的工具包，里面有各种各样的模型以及各种数据集对象和对于数据进行transform的函数，我们从torchvision中导入已有的resnet18

之后声明一个writer对象

writer = SummaryWriter(log_dir='./log', comment='resnet18')

两个变量，分别表示事件存放的文件夹，以及comment表示事件的title

最后在writer内add graph，至于为什么要with writer，试了一下直接如下写，不work

writer.add_graph(net, (input_data,))  # 这种直接的方式并不work

我想的是，可能是因为需要inference中间的节点的data shape才要写成with的吧

 

仍然是运行tensorboard，在浏览器中打开

支持鼠标滚轮放大缩小，拖动，双击可以查看更细节的网络结构，而且数据流箭头上有数据的shape，使用起来非常方便