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，使用起來非常友善