pytorch+tensorboard可視化最簡單例子

前言

儘管pytorch 已經集成了tensorboard的接口，可是你還要下載安裝tensorboard工具。

下載tensorboard：

pip install tensorboard.   

不行的話，再安裝tensorboardX，是早些時候專門給pytorch用的tensorboard。

pip install tensorboardX。

效果

tensorboard用網頁的方式把不少的信息都展示出來，比較方便。上方image和graph分別表明你訓練的數據和你的深度學習網絡結構圖。

最簡單的例子講解

定義一個學習網絡，來分類FashionMNIST，在SummaryWriter的時候，就開始用tensorboard了。
我會分段講解，可是最好是先在文末拷貝總體代碼再回來對照代碼看。

首先import，和定義一些工具類，沒什麼好說的。get_num_correct函數是獲得預測結果和label相同數目的函數。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
 
import torchvision
import torchvision.transforms as transforms
 
from torch.utils.tensorboard import SummaryWriter

def get_num_correct(preds,labels):
    return preds.argmax(dim=1).eq(labels).sum().item()

定義網絡

class Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(in_channels=1,out_channels=6,kernel_size=5)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
        self.fc1=nn.Linear(in_features=12*4*4,out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=60)
        self.out = nn.Linear(in_features=60, out_features=10)
    def forward(self, t):
        t=F.relu(self.conv1(t))
        t=F.max_pool2d(t,kernel_size=2,stride=2)
 
        t = F.relu(self.conv2(t))
        t = F.max_pool2d(t,kernel_size=2,stride=2)
 
        t=t.flatten(start_dim=1)
        t=F.relu(self.fc1(t))
 
        t=F.relu(self.fc2(t))
        t=self.out(t)
 
        return t

main函數裏面，經過pytorch的工具類torchvision導入MNIST數據集，而後用data loader加載進來，爲訓練作準備。

if __name__ == '__main__':
    train_set=torchvision.datasets.FashionMNIST(
        root='./data-source',
        train=True,
        download=True,
        transform=transforms.Compose([
            transforms.ToTensor()
        ])
    )
 
    train_loader=torch.utils.data.DataLoader(train_set,batch_size=100,shuffle=True)

(續上main函數)接着聲明的summary writer就是用到tensorboard的類，tensorboard可以記錄模型學過程當中的不少量，而後用圖表的方式顯示出來。

#tensor board
    tb=SummaryWriter()
    network=Network()
#取出訓練用圖
    images,labels=next(iter(train_loader))
    grid=torchvision.utils.make_grid(images)
#想用tensorboard看什麼，你就tb.add什麼。image、graph、scalar等
    tb.add_image('images', grid)
    tb.add_graph(model=network,input_to_model=images)
    tb.close()
    exit(0)

寫好代碼以後，運行一遍，看有沒有錯誤，有錯誤的地方tensorboard不會儲存也不會顯示。

運行以後這個目錄下會出現runs目錄，裏面儲存量tensorboard要顯示的數據。

而後在這個目錄下cmd，指定吧runs目錄下的數據在tensorboard顯示，開啓tensorboard服務

tensorboard --logdir=runs

而後會出現這個

這樣在瀏覽器訪問本地服務6006端口就能夠看到開頭的效果了。

最後，完整代碼

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
 
import torchvision
import torchvision.transforms as transforms
 
from torch.utils.tensorboard import SummaryWriter
 
def get_num_correct(preds,labels):
    return preds.argmax(dim=1).eq(labels).sum().item()
 
class Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(in_channels=1,out_channels=6,kernel_size=5)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
        self.fc1=nn.Linear(in_features=12*4*4,out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=60)
        self.out = nn.Linear(in_features=60, out_features=10)
    def forward(self, t):
        t=F.relu(self.conv1(t))
        t=F.max_pool2d(t,kernel_size=2,stride=2)
 
        t = F.relu(self.conv2(t))
        t = F.max_pool2d(t,kernel_size=2,stride=2)
 
        t=t.flatten(start_dim=1)
        t=F.relu(self.fc1(t))
 
        t=F.relu(self.fc2(t))
        t=self.out(t)
 
        return t
 
if __name__ == '__main__':
    train_set=torchvision.datasets.FashionMNIST(
        root='./data-source',
        train=True,
        download=True,
        transform=transforms.Compose([
            transforms.ToTensor()
        ])
    )
 
    train_loader=torch.utils.data.DataLoader(train_set,batch_size=100,shuffle=True)
 
    #tensor board
    tb=SummaryWriter()
    network=Network()
#取出訓練用圖
    images,labels=next(iter(train_loader))
    grid=torchvision.utils.make_grid(images)
#想用tensorboard看什麼，你就tb.add什麼。image、graph、scalar等
    tb.add_image('images', grid)
    tb.add_graph(model=network,input_to_model=images)
    tb.close()
    exit(0)