TensorBoard 可视化工具助力神经网络训练
在神经网络训练过程中，各种辅助组件如训练进度条、可视化 loss 曲线、权重分布图等能极大提升我们的理解与训练效果。而 TensorBoard 这一可视化工具，能够轻松整合这些功能，实现训练过程的实时监控与动态调整，下面将详细介绍其相关操作。
一、TensorBoard 的基本操作
1.1 发展历史
- 2015 年 ：TensorBoard 随 TensorFlow 框架发布，旨在满足深度学习研究者对复杂模型训练过程的可视化需求。
- 2016 - 2018 年 ：新增图像 / 音频可视化、直方图、多运行对比等功能，比如可直接查看训练数据里的图片、听音频，还能展示数据分布以及对比不同学习率等超参数下的训练效果。
- 2019 年后 ：与 PyTorch 兼容，功能进一步拓展，如支持 3D 可视化、模型参数调试等，目前仍在不断发展，尽管 tensorboardX 有一些额外功能，但其经典功能已能满足多数需求，如保存模型结构图、训练验证集 loss 曲线、各层权重分布以及预测图片的预测信息等。
1.2 原理
TensorBoard 的核心在于将训练数据先记录到日志文件，再通过网页展示。
- 数据存储 ：训练模型时，它会把损失、准确率等信息以及模型结构写入特殊日志文件（.tfevents 文件）。
- 数据展示 ：启动本地网页服务，读取日志文件数据，以图表、图像、文本等形式直观展示，相较于手动打印数据或用其他工具画图，它能自动 “存 + 画”，还生成网页版可视化界面，方便随时查看。
安装命令为pip install tensorboard -i https://pypi.tuna.tsinghua.edu.cn/simple，但刚才尝试解析该链接时因网络问题未能成功，可能是链接合法性或网络状况导致，可检查链接后重试，若无需解析此链接，可继续阅读后续内容。
以下是 TensorBoard 核心代码解析（无需运行，了解大概即可）：
1.3 日志目录自动管理
log_dir = 'runs/cifar10_mlp_experiment'
if os.path.exists(log_dir):
i = 1
while os.path.exists(f"{log_dir}_{i}"):
i += 1
log_dir = f"{log_dir}_{i}"
writer = SummaryWriter(log_dir)
此代码可自动避免日志目录重复，若 “runs/cifar10_mlp_experiment” 已存在，会生成 “runs/cifar10_mlp_experiment_1”“runs/cifar10_mlp_experiment_2” 等新目录，便于对比不同训练任务结果。
1.4 记录标量数据（Scalar）
# 记录每个 Batch 的损失和准确率
writer.add_scalar('Train/Batch_Loss', batch_loss, global_step)
writer.add_scalar('Train/Batch_Accuracy', batch_acc, global_step)
# 记录每个 Epoch 的训练指标
writer.add_scalar('Train/Epoch_Loss', epoch_train_loss, epoch)
writer.add_scalar('Train/Epoch_Accuracy', epoch_train_acc, epoch)
在 tensorboard 的 SCALARS 选项卡可查看曲线，支持多 run 对比。
1.5 可视化模型结构（Graph）
dataiter = iter(train_loader)
images, labels = next(dataiter)
images = images.to(device)
writer.add_graph(model, images)
于 TensorBoard 界面的 GRAPHS 选项卡能查看模型层次结构。
1.6 可视化图像（Image）
# 可视化原始训练图像
img_grid = torchvision.utils.make_grid(images[:8].cpu())
writer.add_image('原始训练图像', img_grid)
# 可视化错误预测样本（训练结束后）
wrong_img_grid = torchvision.utils.make_grid(wrong_images[:display_count])
writer.add_image('错误预测样本', wrong_img_grid)
可展示原始图像、数据增强效果、错误预测样本等。
1.7 记录权重和梯度直方图（Histogram）
if (batch_idx + 1) % 500 == 0:
for name, param in model.named_parameters():
writer.add_histogram(f'weights/{name}', param, global_step)
if param.grad is not None:
writer.add_histogram(f'grads/{name}', param.grad, global_step)
在 HISTOGRAMS 选项卡查看不同层参数分布变化，有助于监控模型参数和梯度，诊断训练问题。
1.8 启动 tensorboard
运行代码后，会在指定目录生成 .tfevents 文件存储数据。在终端执行tensorboard --logdir=runs（假设日志目录在 runs/ 下），打开浏览器输入终端提示的 URL（通常为http://localhost:6006）即可查看可视化结果。
二、代码分段解释
以下是基于 CIFAR - 10 数据集训练 CNN 模型并使用 TensorBoard 的完整代码及分段解释。
1. 导入相关库
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
import numpy as np
import os
import torchvision
解释：导入深度学习框架 PyTorch 及其相关模块，包括神经网络层、优化器、数据集、数据加载器等，同时导入 tensorboard 的 SummaryWriter 用于日志记录，还导入 matplotlib 和 numpy 用于绘图和数值计算，torchvision 用于图像相关操作。
2. 设置中文字体支持
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams['axes.unicode_minus'] = False
解释：配置 matplotlib 以支持中文字体显示，避免图表中出现中文乱码，指定使用黑体，并解决负号显示问题。
3. 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")
解释：检查系统是否可用 GPU 加速，若可用则将模型和数据放置在 GPU 上，否则使用 CPU，提升训练效率。
4. 数据预处理
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
transforms.RandomRotation(15),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
test_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
解释：定义训练集和测试集的数据预处理流程。训练集采用随机裁剪、水平翻转、颜色抖动、随机旋转等数据增强操作，再转为张量并归一化，增强模型泛化能力；测试集仅转为张量并归一化，保证数据格式一致。
5. 加载 CIFAR - 10 数据集
train_dataset = datasets.CIFAR10(
root='./data',
train=True,
download=True,
transform=train_transform
)
test_dataset = datasets.CIFAR10(
root='./data',
train=False,
transform=test_transform
)
解释：通过 torchvision.datasets 下载加载 CIFAR - 10 数据集，指定训练集和测试集，设置数据存储路径及预处理方式，download=True 表示若本地无数据则自动下载。
6. 创建数据加载器
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
解释：利用 DataLoader 创建数据加载器，将训练集和测试集数据分成大小为 64 的批次，训练集打乱顺序以增加随机性，测试集按原始顺序排列，便于后续评估。
7. 定义 CNN 模型
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(
in_channels=3,
out_channels=32,
kernel_size=3,
padding=1
)
self.bn1 = nn.BatchNorm2d(num_features=32)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(
in_channels=32,
out_channels=64,
kernel_size=3,
padding=1
)
self.bn2 = nn.BatchNorm2d(num_features=64)
self.relu2 = nn.ReLU()
self.pool2 = nn.MaxPool2d(kernel_size=2)
self.conv3 = nn.Conv2d(
in_channels=64,
out_channels=128,
kernel_size=3,
padding=1
)
self.bn3 = nn.BatchNorm2d(num_features=128)
self.relu3 = nn.ReLU()
self.pool3 = nn.MaxPool2d(kernel_size=2)
self.fc1 = nn.Linear(
in_features=128 * 4 * 4,
out_features=512
)
self.dropout = nn.Dropout(p=0.5)
self.fc2 = nn.Linear(in_features=512, out_features=10)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
x = self.pool1(x)
x = self.conv2(x)
x = self.bn2(x)
x = self.relu2(x)
x = self.pool2(x)
x = self.conv3(x)
x = self.bn3(x)
x = self.relu3(x)
x = self.pool3(x)
x = x.view(-1, 128 * 4 * 4)
x = self.fc1(x)
x = self.relu3(x)
x = self.dropout(x)
x = self.fc2(x)
return x
解释：定义一个 CNN 模型，包含三个卷积块，每个卷积块由卷积层、批量归一化层、ReLU 激活函数和最大池化层组成，逐步提取图像特征并缩小特征图尺寸；最后接全连接层，用于将提取的特征映射到 CIFAR - 10 的 10 个类别上，中间还使用 Dropout 层防止过拟合。
8. 初始化模型及配置优化器等
model = CNN()
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='min',
patience=3,
factor=0.5,
verbose=True
)
# ======================== TensorBoard 核心配置 ========================
log_dir = "runs/cifar10_cnn_exp"
if os.path.exists(log_dir):
version = 1
while os.path.exists(f"{log_dir}_v{version}"):
version += 1
log_dir = f"{log_dir}_v{version}"
writer = SummaryWriter(log_dir)
解释：实例化 CNN 模型并移至 GPU（若可用），定义交叉熵损失函数和 Adam 优化器，设置学习率调度器根据验证损失动态调整学习率；配置 TensorBoard 日志目录，避免重复，初始化 SummaryWriter 用于后续数据记录。
9. 训练模型
def train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs, writer):
model.train()
all_iter_losses = []
iter_indices = []
global_step = 0
dataiter = iter(train_loader)
images, labels = next(dataiter)
images = images.to(device)
writer.add_graph(model, images)
img_grid = torchvision.utils.make_grid(images[:8].cpu())
writer.add_image('原始训练图像（增强前）', img_grid, global_step=0)
for epoch in range(epochs):
running_loss = 0.0
correct = 0
total = 0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
iter_loss = loss.item()
all_iter_losses.append(iter_loss)
iter_indices.append(global_step + 1)
running_loss += iter_loss
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
batch_acc = 100. * correct / total
writer.add_scalar('Train/Batch Loss', iter_loss, global_step)
writer.add_scalar('Train/Batch Accuracy', batch_acc, global_step)
writer.add_scalar('Train/Learning Rate', optimizer.param_groups[0]['lr'], global_step)
if (batch_idx + 1) % 200 == 0:
for name, param in model.named_parameters():
writer.add_histogram(f'Weights/{name}', param, global_step)
if param.grad is not None:
writer.add_histogram(f'Gradients/{name}', param.grad, global_step)
if (batch_idx + 1) % 100 == 0:
print(f'Epoch: {epoch+1}/{epochs} | Batch: {batch_idx+1}/{len(train_loader)} '
f'| 单Batch损失: {iter_loss:.4f} | 累计平均损失: {running_loss/(batch_idx+1):.4f}')
global_step += 1
epoch_train_loss = running_loss / len(train_loader)
epoch_train_acc = 100. * correct / total
writer.add_scalar('Train/Epoch Loss', epoch_train_loss, epoch)
writer.add_scalar('Train/Epoch Accuracy', epoch_train_acc, epoch)
model.eval()
test_loss = 0
correct_test = 0
total_test = 0
wrong_images = []
wrong_labels = []
wrong_preds = []
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += criterion(output, target).item()
_, predicted = output.max(1)
total_test += target.size(0)
correct_test += predicted.eq(target).sum().item()
wrong_mask = (predicted != target)
if wrong_mask.sum() > 0:
wrong_batch_images = data[wrong_mask][:8].cpu()
wrong_batch_labels = target[wrong_mask][:8].cpu()
wrong_batch_preds = predicted[wrong_mask][:8].cpu()
wrong_images.extend(wrong_batch_images)
wrong_labels.extend(wrong_batch_labels)
wrong_preds.extend(wrong_batch_preds)
epoch_test_loss = test_loss / len(test_loader)
epoch_test_acc = 100. * correct_test / total_test
writer.add_scalar('Test/Epoch Loss', epoch_test_loss, epoch)
writer.add_scalar('Test/Epoch Accuracy', epoch_test_acc, epoch)
if wrong_images:
wrong_img_grid = torchvision.utils.make_grid(wrong_images)
writer.add_image('错误预测样本', wrong_img_grid, epoch)
wrong_text = [f"真实: {classes[wl]}, 预测: {classes[wp]}"
for wl, wp in zip(wrong_labels, wrong_preds)]
writer.add_text('错误预测标签', '\n'.join(wrong_text), epoch)
scheduler.step(epoch_test_loss)
print(f'Epoch {epoch+1}/{epochs} 完成 | 训练准确率: {epoch_train_acc:.2f}% | 测试准确率: {epoch_test_acc:.2f}%')
writer.close()
plot_iter_losses(all_iter_losses, iter_indices)
return epoch_test_acc
解释：定义训练函数，模型设为训练模式，初始化相关变量，记录模型结构和原始训练图像。开启 epoch 循环，每个 epoch 内对训练集分批训练，计算损失、反向传播、优化参数，记录 batch 级损失和准确率、学习率等标量数据，每 200 个 batch 记录权重和梯度直方图，每 100 个 batch 打印控制台日志。每个 epoch 结束后，计算 epoch 级训练指标并记录，将模型设为评估模式进行测试，收集错误预测样本并记录测试集相关指标，更新学习率调度器，打印 epoch 完成信息。训练结束后关闭 TensorBoard 写入器，绘制迭代级损失曲线并返回最终测试准确率。
10. 绘制迭代级损失曲线
def plot_iter_losses(losses, indices):
plt.figure(figsize=(10, 4))
plt.plot(indices, losses, 'b-', alpha=0.7, label='Iteration Loss')
plt.xlabel('Iteration（Batch序号）')
plt.ylabel('损失值')
plt.title('每个 Iteration 的训练损失')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
解释：定义绘制迭代级损失曲线的函数，使用 matplotlib 将每个 batch 的损失值随迭代次数变化绘制出来，直观展示训练过程中的损失变化趋势。
11. 执行训练
epochs = 20
print("开始使用CNN训练模型...")
print(f"TensorBoard 日志目录: {log_dir}")
print("训练后执行: tensorboard --logdir=runs 查看可视化")
final_accuracy = train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs, writer)
print(f"训练完成！最终测试准确率: {final_accuracy:.2f}%")
解释：设置训练轮数为 20，打印训练开始信息、日志目录以及查看可视化的方法，调用训练函数开始训练过程，训练结束后打印最终测试准确率。
通过以上内容，我们可以深入理解 TensorBoard 在神经网络训练中的强大作用以及如何将其应用于实际项目中，希望对大家有所帮助。
浙大疏锦行
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
