tensorboard查看pytorch模型

與tensorflow模型與caffe模型不同，當前的pytorch沒有官方的直觀查看網絡結構的工具，google了下pytorch的網絡解析的方法，發現可以將pytorch的model轉換成為events文件使用tensorboard查看，記錄之。

安裝插件

• TensorboardX，TensorboardX支持scalar, image, figure, histogram, audio, text, graph, onnx_graph, embedding, pr_curve and videosummaries等不同的可視化展示方式，具體介紹移步至項目Github 觀看詳情。使用下面的命令安裝

  pip install tensorboardX
  

• 安裝tensorboard，參考命令

  pip install tensorboard
  

具體過程

參考代碼

#-*-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,))

該代碼中14行聲明一個writer對象，分別表示events存放的目錄，comment表示事件的title，然后使用如下的方式打開tensorboard

tensorboard --logpath=D:\log --port=6006

然后按照命令行提示打開即可。

參考鏈接

• PyTorch 使用 TensorboardX 進行網絡可視化-PyTorch 中文網
• Pytorch使用tensorboardX可視化。超詳細！！！ - 簡書
• torch.utils.tensorboard — PyTorch master documentation