TensorBoard:训练日志及网络结构可视化工具

1.安装

pip install tensorboard

2.可视化标量数据(loss和accuracy)

from torch.utils.tensorboard import SummaryWriter
import random

# 实例化TensorBoard
logs_writer = SummaryWriter('./logs')

# 可视化标量数据
for epoch_id in range(100):
    logs_writer.add_scalar("train/loss",random.random(),epoch_id)
    logs_writer.add_scalar("train/accuracy",random.random(),epoch_id)
    
    logs_writer.add_scalar("test/loss",random.random(),epoch_id)
    logs_writer.add_scalar("test/accuracy",random.random(),epoch_id)

tensorboard --logdir=logs --bind_all

3.可视化网络结构

"""
可视化网络结构
"""

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter

# 实例化TensorBoard
writer = SummaryWriter('./model_vis')


# 构建网络模型 - 使用自定义类
class Digit_Rec(nn.Module):
    def __init__(self):
        super(Digit_Rec,self).__init__()
        self.conv1 = nn.Conv2d(1,10,5) #1：灰度图片的通道，10：输出通道，5：kernel
        self.relu1 = nn.ReLU()
        self.max_pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(10,20,3) #10：输入通道，20：输出通道，3：Kernel
        self.relu2 = nn.ReLU()
        self.fc1 = nn.Linear(20*10*10,500) # 20*10*10：输入通道，500：输出通道
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(500,10) # 500：输入通道，10：输出通道
        self.relu4 = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
    def forward(self,x):
        batch_size = x.size(0) # x的格式：batch_size x 1 x 28 x 28 拿到了batch_size
        
        x = self.conv1(x) # 输入：batch*1*28*28   输出：batch*10*24*24
        x = self.relu1(x)
        
        x = self.max_pool(x) # 输入：batch*10*24*24输出：batch*10*12*12
        
        x = self.conv2(x)
        x = self.relu2(x)
        
        x = x.view(batch_size,-1) #fatten 展平 -1自动计算维度，20*10*10=2000
        
        x = self.fc1(x) # 输入：batch*2000  输出：batch*500
        x = self.relu3(x)
        
        x = self.fc2(x) # 输入：batch*500  输出：batch*10
        x = self.relu4(x)
        
        output = self.softmax(x) # 计算分类后，每个数字的概率值
        
        return output

model =  Digit_Rec()

model = Digit_Rec()
images = torch.randn(1, 1, 28, 28)
writer.add_graph(model, images)
writer.close()

tensorboard --logdir=model_vis --bind_all

参考资料

1. Pytorch中使用tensorboard学习笔记(2)记录损失loss和准确率accuracy
2. pytorch中使用tensorboard绘制Accuracy/Loss曲线（train和test显示在同一幅图中）
3. pytorch中使用TensorBoard进行可视化Loss及特征图