pytorch使用TensorBoard可视化损失函数曲线、精度信息

在训练神经网络时，我们希望能够直观地训练情况，例如损失函数的曲线、输入的图像、模型精度等信息，这些信息可以帮助我们更好地监督网络的训练过程，并为参数优化提供方向和依据。

其实我们可以有个更容易地实现方式就是定义一个列表，然后将每个epoch的训练结果添加到列表中，待模型训练完成之后，使用这个列表中的数据进行可视化以及绘图操作等。

本文提供一个更为专业的操作，它是一个常用的可视化工具：TensorBoard

PyTorch已经内置了TensorBoard的相关接口，用户在安装后便可调用相关接口进行数据可视化。

下面我将使用LeNet训练MNIST为例，来讲解tensorboard如何使用。

加载数据

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms as T
from tqdm import tqdm

data_transform = T.Compose([
        T.RandomResizedCrop(32),
        T.ToTensor(),
])

train_dataset = torchvision.datasets.MNIST('./', train=True, transform=data_transform)
val_dataset = torchvision.datasets.MNIST('./', train=False, transform=data_transform)

train_loader = torch.utils.data.DataLoader(train_dataset, 32)
val_loader = torch.utils.data.DataLoader(val_dataset, 32)

定义网络

class LeNet5(nn.Module):
    def __init__(self,num_class=10):
        super(LeNet5,self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.pool1 = nn.AvgPool2d((2, 2))
        
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.pool2 = nn.AvgPool2d((2, 2))
        
        self.conv3 = nn.Conv2d(16, 120, 5)
        
        self.relu = nn.ReLU()

        self.fc1 = nn.Linear(120, 84)
        self.fc2 = nn.Linear(84, num_class)
        
    def forward(self, x):
        # x: torch.Size([1, 3, 32, 32])
        
        x = self.conv1(x) # torch.Size([1, 6, 28, 28])
        x = self.relu(x)
        x = self.pool1(x) # torch.Size([1, 6, 14, 14])

        x = self.conv2(x) # torch.Size([1, 16, 10, 10])
        x = self.relu(x)
        x = self.pool2(x) # torch.Size([1, 16, 5, 5])

        x = self.conv3(x) # torch.Size([1, 120, 1, 1])
        x = self.relu(x)
        
        x = x.flatten(start_dim=1) # torch.Size([1, 120])
        
        x = self.fc1(x) # torch.Size([1, 84])
        x = self.relu(x)
        x = self.fc2(x) # torch.Size([1, 5])

        return x

定义模型等组件

model = LeNet5(10)
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), 0.003)

# 日志信息保存路径
logger = SummaryWriter(log_dir='./log')

模型训练

for epoch in range(50):
    train_accuracy = 0
    train_loss = 0
    val_accuracy = 0
    val_loss = 0
    epoch_acc_count = 0
    count = 0
    running_loss = 0
    
    model.train()
    for data in tqdm(train_loader):
        images, labels = data
        optimizer.zero_grad()
        output = model(images)
        loss = loss_function(output, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()

        # 计算每个epoch正确的个数
        epoch_acc_count += (output.argmax(axis=1) == labels.view(-1)).sum()
        count += len(images)
    
    # 写入日志信息
    logger.add_scalar('train_loss', running_loss, epoch + 1)
    logger.add_scalar('train_accuracy', epoch_acc_count / count, epoch + 1)
    
    running_loss = 0
    epoch_acc_count = 0
    count = 0
    
    model.eval()
    for data in tqdm(val_loader):
        images, labels = data
        optimizer.zero_grad()
        output = model(images)
        loss = loss_function(output, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()

        # 计算每个epoch正确的个数
        epoch_acc_count += (output.argmax(axis=1) == labels.view(-1)).sum()
        count += len(images)
    
    logger.add_scalar('val_loss', running_loss, epoch + 1)
    logger.add_scalar('val_accuracy', epoch_acc_count / count, epoch + 1)

命令启动TensorBoard

通过下面的命令行参数即可启动tensorboard，log_path就是我们模型训练时日志保存的路径。

tensorboard --logdir=log_path