我們知道,對於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,使用起來非常方便