PyTorch训练监控工具：使用TensorBoard跟踪性能指标

在PyTorch深度学习项目中，有效的训练监控是优化模型性能的关键环节。本文将展示如何使用TensorBoard来跟踪PyTorch模型的训练指标，并提供具体的可复现代码示例。

环境准备

首先安装必要的依赖包：

pip install torch torchvision tensorboard

基础监控实现

创建一个简单的训练循环并在TensorBoard中记录指标：

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

# 初始化TensorBoard写入器
writer = SummaryWriter('runs/mnist_training')

# 模型定义
model = nn.Sequential(
    nn.Conv2d(1, 32, 3, 1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(512, 10)
)

# 优化器和损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 训练循环
for epoch in range(5):
    running_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        pred = output.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
        
    # 记录每个epoch的平均损失和准确率
    avg_loss = running_loss / len(train_loader)
    accuracy = 100. * correct / total
    
    writer.add_scalar('Loss/Train', avg_loss, epoch)
    writer.add_scalar('Accuracy/Train', accuracy, epoch)
    
    print(f'Epoch {epoch}: Loss={avg_loss:.4f}, Accuracy={accuracy:.2f}%')

writer.close()

高级监控功能

除了基本的损失和准确率，还可以记录更多实用指标：

# 记录学习率
writer.add_scalar('Params/Learning_Rate', optimizer.param_groups[0]['lr'], epoch)

# 记录参数梯度分布
for name, param in model.named_parameters():
    if param.grad is not None:
        writer.add_histogram(f'Gradients/{name}', param.grad, epoch)

# 记录模型权重分布
for name, param in model.named_parameters():
    writer.add_histogram(f'Weights/{name}', param, epoch)

性能测试数据

在MNIST数据集上进行测试，使用以下配置：

• 模型：简单CNN结构
• 优化器：Adam(lr=0.001)
• 批次大小：64
• 训练轮数：5轮

通过TensorBoard可视化界面，可以实时监控这些指标的变化趋势，快速识别模型训练中的问题。

启动TensorBoard

# 在项目根目录运行
tensorboard --logdir=runs