PyTorch 使用 TensorBoard 可视化模型，数据和训练

在深度学习模型的训练过程中，可视化是一个不可或缺的环节。它可以帮助我们更好地理解模型的行为、数据的分布以及训练过程的进展。PyTorch 与 TensorBoard 的深度集成，为我们提供了强大的可视化工具。本教程将指导你如何使用 PyTorch 和 TensorBoard 进行模型训练的可视化，助力你在深度学习的道路上事半功倍。

一、环境准备

在开始之前，确保你已经安装了以下必要的库：

import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter

二、数据加载与预处理

我们将使用 Fashion-MNIST 数据集，这是一个广泛用于分类任务的数据集，包含 10 个类别的服装图像。

## 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])


## 加载数据集
trainset = torchvision.datasets.FashionMNIST('./data', download=True, train=True, transform=transform)
testset = torchvision.datasets.FashionMNIST('./data', download=True, train=False, transform=transform)


## 创建数据加载器
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)


## 定义类别名称
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')

三、定义模型架构

我们定义一个简单的卷积神经网络，用于 Fashion-MNIST 数据集的分类任务。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)  # 输入通道1，输出通道6，卷积核大小5
        self.pool = nn.MaxPool2d(2, 2)   # 最大池化层，窗口大小2，步长2
        self.conv2 = nn.Conv2d(6, 16, 5) # 输入通道6，输出通道16，卷积核大小5
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)


    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # 卷积 + 激活 + 池化
        x = self.pool(F.relu(self.conv2(x)))  # 卷积 + 激活 + 池化
        x = x.view(-1, 16 * 4 * 4)            # 展平操作
        x = F.relu(self.fc1(x))               # 全连接 + 激活
        x = F.relu(self.fc2(x))               # 全连接 + 激活
        x = self.fc3(x)                       # 输出层
        return x


net = Net()
print(net)

四、设置 TensorBoard

创建一个 SummaryWriter 对象，用于将信息写入 TensorBoard。

## 设置 TensorBoard
writer = SummaryWriter('runs/fashion_mnist_experiment')

五、写入数据到 TensorBoard

将训练数据和模型结构写入 TensorBoard，以便进行可视化。

## 获取一批训练数据
dataiter = iter(trainloader)
images, labels = next(dataiter)


## 创建图像网格
img_grid = torchvision.utils.make_grid(images)


## 将图像写入 TensorBoard
writer.add_image('Fashion-MNIST_images', img_grid)


## 将模型结构写入 TensorBoard
writer.add_graph(net, images)


## 关闭 TensorBoard writer
writer.close()

六、训练模型并跟踪训练过程

在训练过程中，我们将损失值和模型预测结果写入 TensorBoard，以便实时跟踪训练进度。

## 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)


## 训练模型
running_loss = 0.0
for epoch in range(2):  # 遍历数据集多次
    for i, data in enumerate(trainloader, 0):
        # 获取输入数据和标签
        inputs, labels = data


        # 清空梯度缓存
        optimizer.zero_grad()


        # 前向传播 + 反向传播 + 优化
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()


        # 累计损失
        running_loss += loss.item()


        # 每 1000 个批次写入 TensorBoard
        if i % 1000 == 999:
            writer.add_scalar('training_loss', running_loss / 1000, epoch * len(trainloader) + i)
            running_loss = 0.0


print('Finished Training')

七、评估模型并生成可视化报告

在测试集上评估模型性能，并生成精度召回曲线等可视化报告。

## 在测试集上评估模型
class_probs = []
class_preds = []
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        probs = torch.softmax(outputs, dim=1)
        _, preds = torch.max(outputs, 1)
        class_probs.append(probs)
        class_preds.append(preds)


test_probs = torch.cat(class_probs)
test_preds = torch.cat(class_preds)


## 为每个类别生成精度召回曲线
for i in range(len(classes)):
    writer.add_pr_curve(classes[i], test_preds == i, test_probs[:, i])


## 关闭 TensorBoard writer
writer.close()

八、总结

通过本教程，你已经掌握了如何使用 PyTorch 和 TensorBoard 进行深度学习模型的可视化。 TensorBoard 提供了丰富的可视化工具，可以帮助你更好地理解模型的行为和训练过程。希望你在编程狮的学习平台上能够充分利用这些工具，提升你的深度学习项目开发能力！

在编程狮平台学习的你，现在可以尝试以下实践项目来加深对 TensorBoard 可视化的理解：

• 可视化自定义数据集 ：使用你在编程狮平台上学习到的自定义数据集知识，结合 TensorBoard 的可视化功能，展示不同数据增强方法对图像数据的影响。例如，对比原始图像与经过旋转、裁剪、翻转等变换后的图像，观察这些变换如何提高模型的泛化能力。
• 跟踪模型性能指标 ：在训练一个更复杂的模型（如用于图像分割或自然语言处理的模型）时，使用 TensorBoard 记录多个性能指标，如准确率、召回率、F1 值等。分析这些指标在训练过程中的变化趋势，找出模型优化的方向。