pytorch使用tensorboardX进行网络可视化

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

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

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

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

先贴上代码

 1 #-*-coding:utf-8-*-
 2 import torch
 3 import torchvision
 4 from torch.autograd import Variable
 5 from tensorboardX import SummaryWriter
 6 
 7 # 模拟输入数据
 8 input_data = Variable(torch.rand(16, 3, 224, 224))
 9 
10 # 从torchvision中导入已有模型
11 net = torchvision.models.resnet18()
12 
13 # 声明writer对象，保存的文件夹，异己名称
14 writer = SummaryWriter(log_dir='./log', comment='resnet18')
15 with writer:
16     writer.add_graph(net, (input_data,))

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

以后声明一个writer对象

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

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

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

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

我想的是，多是由于须要inference中间的节点的data shape才要写成with的吧

 

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

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