使用 PyTorch 和 TensorBoard 实时可视化模型训练

在这个教程中，我们将使用 PyTorch 训练一个简单的多层感知机（MLP）模型来解决 MNIST 手写数字分类问题，并且使用 TensorBoard 来可视化训练过程中的不同信息，如损失、准确度、图像、参数分布和学习率变化。

---

步骤 1：安装必要的库

首先，确保你已经安装了必要的 Python 库。如果尚未安装，可以通过以下命令进行安装：

pip install torch torchvision tensorboard

torch 和 torchvision 是 PyTorch 中的核心库，而 tensorboard 是用于在浏览器中可视化训练过程的工具。

---

步骤 2：导入库

导入 PyTorch 及相关的库：

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

这些库分别用于：

• torch：PyTorch 核心库，包含各种模型和优化工具。
• torchvision：包含数据集和预训练模型。
• SummaryWriter：用于将训练过程中的数据写入 TensorBoard。
• os：文件系统操作，用于创建日志目录。

---

步骤 3：加载和预处理数据

我们将使用 MNIST 数据集，它包含手写数字的图像。使用 torchvision 来加载并应用预处理：

# 数据预处理
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))  # 归一化
])# 下载训练集和测试集
train_set = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)test_set = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=1000, shuffle=False)

• transforms.ToTensor()：将图像数据从 PIL 图像转换为 PyTorch 张量。
• transforms.Normalize()：对图像进行归一化处理，将数据标准化为均值为 0，标准差为 1。

train_loader 和 test_loader 分别是用于训练和测试的 PyTorch 数据加载器。

---

步骤 4：定义模型

我们将创建一个简单的 多层感知机（MLP），包含两个全连接层，使用 ReLU 激活函数。

# 简单的多层感知机 (MLP)
class SimpleMLP(nn.Module):def __init__(self):super(SimpleMLP, self).__init__()self.fc1 = nn.Linear(28*28, 128)self.fc2 = nn.Linear(128, 64)self.fc3 = nn.Linear(64, 10)  # 10个输出类def forward(self, x):x = x.view(-1, 28*28)  # 展平图像x = F.relu(self.fc1(x))  # 激活函数x = F.relu(self.fc2(x))  # 激活函数return self.fc3(x)  # 输出层

fc1, fc2, fc3 分别是模型的全连接层，每一层的输入和输出维度通过 Linear 层来定义。

---

步骤 5：设置损失函数和优化器

选择适当的损失函数和优化器来训练模型：

# 实例化模型
model = SimpleMLP()# 使用 GPU 或 CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

• CrossEntropyLoss 是用于分类问题的标准损失函数。
• Adam 优化器是一种自适应学习率的优化方法，适用于大多数任务。

---

步骤 6：设置 TensorBoard Writer

为 TensorBoard 创建一个日志目录，并初始化 SummaryWriter：

# TensorBoard 日志保存路径
log_dir = "./runs/mnist_example"
if not os.path.exists(log_dir):os.makedirs(log_dir)# 初始化 TensorBoard Writer
writer = SummaryWriter(log_dir)# 添加模型的计算图 (graph)
sample_data, _ = next(iter(train_loader))
writer.add_graph(model, sample_data.to(device))

• log_dir 是 TensorBoard 保存日志文件的目录。
• SummaryWriter 用于记录训练过程中的数据。
• add_graph 用于将模型结构写入 TensorBoard。

---

步骤 7：训练模型

接下来，我们定义训练过程。在每 100 步记录一次损失和准确度，每 1000 步记录一次训练图像。

# 训练周期
epochs = 10
global_step = 0
for epoch in range(epochs):running_loss = 0.0correct = 0total = 0# 训练阶段model.train()for i, (inputs, labels) in enumerate(train_loader):inputs, labels = inputs.to(device), labels.to(device)# 清空梯度optimizer.zero_grad()# 前向传播outputs = model(inputs)# 计算损失loss = criterion(outputs, labels)loss.backward()  # 反向传播optimizer.step()  # 更新参数# 累积损失和精度running_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels).sum().item()# 每 100 步记录一次if (i + 1) % 100 == 0:avg_loss = running_loss / 100acc = correct / totalprint(f"[Epoch {epoch+1}, Step {i+1}] Loss: {avg_loss:.4f} | Accuracy: {acc:.4f}")# 记录损失和准确度writer.add_scalar("Loss/train", avg_loss, global_step)writer.add_scalar("Accuracy/train", acc, global_step)# 记录每层的权重分布for name, param in model.named_parameters():writer.add_histogram(name, param, global_step)running_loss = 0.0correct = 0total = 0# 记录每1000步的图像（展示输入图像）if (i + 1) % 1000 == 0:img_grid = torchvision.utils.make_grid(inputs[:16])  # 取16个样本writer.add_image(f"Train_Images/{epoch+1}_{i+1}", img_grid, global_step)global_step += 1

• writer.add_scalar 用于记录训练中的损失和准确度。
• writer.add_histogram 用于记录每层参数的权重分布。
• writer.add_image 用于记录训练样本的图像。

---

步骤 8：测试模型

测试模型并记录测试集的损失和准确度。

# 测试阶段
model.eval()
test_loss = 0.0
correct = 0
total = 0
with torch.no_grad():for inputs, labels in test_loader:inputs, labels = inputs.to(device), labels.to(device)outputs = model(inputs)loss = criterion(outputs, labels)test_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels).sum().item()# 记录测试集的损失和准确率
test_accuracy = correct / total
test_loss /= len(test_loader)
writer.add_scalar("Loss/test", test_loss, global_step)
writer.add_scalar("Accuracy/test", test_accuracy, global_step)# 测试集的图像展示
test_images, _ = next(iter(test_loader))
img_grid = torchvision.utils.make_grid(test_images[:16])
writer.add_image('Test_Images', img_grid)

---

步骤 9：关闭 TensorBoard Writer

完成训练和测试后，关闭 TensorBoard 写入器：

# 关闭 TensorBoard 写入器
writer.close()

---

步骤 10：启动 TensorBoard

1. 在终端中运行以下命令启动 TensorBoard：

   tensorboard --logdir=./runs/mnist_example
   

2. 打开浏览器，访问 http://localhost:6006 即可查看：

   • 损失和精度曲线
   • 每层权重分布
   • 训练/测试图像样本
   • 学习率的变化

---

总结

在本教程中，我们展示了如何使用 PyTorch 训练一个简单的 MLP 模型，利用 TensorBoard 记录并实时可视化以下内容：

• 损失和准确度曲线
• 每层的权重和梯度分布
• 训练过程中的图像
• 学习率的变化

TensorBoard 提供了一种强大的可视化方式，可以帮助你深入了解模型的训练过程，并有效地进行调优。

好的！以下是 完整的代码，包括每个部分的实现。每个代码块的功能都是按步骤分开的，确保你可以从头到尾顺利运行。

完整代码：PyTorch + TensorBoard

# 导入所需的库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
import os# tensorboard --logdir=./runs/mnist_example# 设置随机种子（可选）
torch.manual_seed(0)# 数据预处理
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))  # 归一化
])# 下载训练集和测试集
train_set = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)test_set = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=1000, shuffle=False)# 定义简单的多层感知机 (MLP) 模型
class SimpleMLP(nn.Module):def __init__(self):super(SimpleMLP, self).__init__()self.fc1 = nn.Linear(28*28, 128)self.fc2 = nn.Linear(128, 64)self.fc3 = nn.Linear(64, 10)  # 10个输出类def forward(self, x):x = x.view(-1, 28*28)  # 展平图像x = F.relu(self.fc1(x))  # 激活函数x = F.relu(self.fc2(x))  # 激活函数return self.fc3(x)  # 输出层# 实例化模型
model = SimpleMLP()# 使用 GPU 或 CPU
# 设置设备并打印详细信息
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if device.type == "cuda":device_name = torch.cuda.get_device_name(device)
else:device_name = "CPU"print("=" * 90)
print(f"Device set to : {device_name}")
print("=" * 90)model.to(device)# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 设置 TensorBoard Writer
log_dir = "./runs/mnist_example"
if not os.path.exists(log_dir):os.makedirs(log_dir)# 初始化 TensorBoard Writer
writer = SummaryWriter(log_dir)# 添加模型的计算图 (graph)
sample_data, _ = next(iter(train_loader))
writer.add_graph(model, sample_data.to(device))# 设置学习率调度器（可选）
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)# 训练模型
epochs = 10
global_step = 0
for epoch in range(epochs):running_loss = 0.0correct = 0total = 0# 训练阶段model.train()for i, (inputs, labels) in enumerate(train_loader):inputs, labels = inputs.to(device), labels.to(device)# 清空梯度optimizer.zero_grad()# 前向传播outputs = model(inputs)# 计算损失loss = criterion(outputs, labels)loss.backward()  # 反向传播optimizer.step()  # 更新参数# 累积损失和精度running_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels).sum().item()# 每 100 步记录一次if (i + 1) % 100 == 0:avg_loss = running_loss / 100acc = correct / totalprint(f"[Epoch {epoch+1}, Step {i+1}] Loss: {avg_loss:.4f} | Accuracy: {acc:.4f}")# 记录损失和准确度writer.add_scalar("Loss/train", avg_loss, global_step)writer.add_scalar("Accuracy/train", acc, global_step)# 记录每层的权重分布for name, param in model.named_parameters():writer.add_histogram(name, param, global_step)running_loss = 0.0correct = 0total = 0# 记录每1000步的图像（展示输入图像）if (i + 1) % 1000 == 0:img_grid = torchvision.utils.make_grid(inputs[:16])  # 取16个样本writer.add_image(f"Train_Images/{epoch+1}_{i+1}", img_grid, global_step)global_step += 1# 更新学习率lr_scheduler.step()# 每个epoch记录学习率变化writer.add_scalar("Learning_Rate", optimizer.param_groups[0]['lr'], epoch)# 测试模型
model.eval()
test_loss = 0.0
correct = 0
total = 0
with torch.no_grad():for inputs, labels in test_loader:inputs, labels = inputs.to(device), labels.to(device)outputs = model(inputs)loss = criterion(outputs, labels)test_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels).sum().item()# 记录测试集的损失和准确率
test_accuracy = correct / total
test_loss /= len(test_loader)
writer.add_scalar("Loss/test", test_loss, global_step)
writer.add_scalar("Accuracy/test", test_accuracy, global_step)# 测试集的图像展示
test_images, _ = next(iter(test_loader))
img_grid = torchvision.utils.make_grid(test_images[:16])
writer.add_image('Test_Images', img_grid)# 关闭 TensorBoard 写入器
writer.close()

---

解释与步骤

1. 导入库：我们首先导入了 torch, torchvision, 和 SummaryWriter 来处理模型的定义、数据加载和 TensorBoard 写入。

2. 数据加载：使用 torchvision.datasets.MNIST 下载并加载 MNIST 数据集，使用 DataLoader 进行批量加载。

3. 模型定义：定义了一个简单的多层感知机（MLP）模型，包括两层全连接层和一个输出层。每层使用 ReLU 激活函数。

4. 损失函数和优化器：使用交叉熵损失函数 (CrossEntropyLoss) 和 Adam 优化器。Adam 是一种常见的自适应优化算法，能有效优化深度学习模型。

5. TensorBoard 设置：

   • 创建一个 SummaryWriter 来将训练过程中的各种信息写入日志。
   • 使用 writer.add_graph(model, sample_data) 记录模型结构。
   • 记录每个训练步骤中的损失、准确度和模型参数（权重）分布。

6. 训练过程：

   • 每 100 步记录一次当前损失和准确度。
   • 每 1000 步展示当前批次的图像样本。

7. 学习率调度器：使用 StepLR 学习率调度器，每训练 5 个 epoch，学习率降低 10 倍。

8. 测试过程：在训练结束后，评估模型在测试集上的表现，并记录测试集的损失和准确度。

9. 关闭 TensorBoard：训练和测试结束后，关闭 SummaryWriter 以保存日志。

---

启动 TensorBoard

1. 运行上面的代码后，终端会生成 TensorBoard 日志。
2. 打开终端，使用以下命令启动 TensorBoard：

tensorboard --logdir=./runs/mnist_example

3. 然后，打开浏览器，访问 http://localhost:6006，即可看到训练过程的可视化界面。

---

在 TensorBoard 中可视化的内容

• Scalars：损失和准确度曲线。
• Histograms：每层权重的分布情况。
• Images：训练图像和测试图像样本。
• Learning Rate：学习率随训练过程的变化。

---