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，使用起來很是方便