不监控 = 盲飞。可视化 = 带上仪表盘

TensorBoard、WandB、tqdm 进度条——全方位训练监控方案。

[toc]

训练监控与可视化


1️⃣ TensorBoard —— PyTorch 内置方案

安装与启动

pip install tensorboard
 
# 启动 TensorBoard 服务
tensorboard --logdir=runs --port=6006
 
# 浏览器打开 http://localhost:6006

基础用法

from torch.utils.tensorboard import SummaryWriter
 
writer = SummaryWriter('runs/experiment_1')
 
for epoch in range(EPOCHS):
    train_loss, train_acc = train_one_epoch(...)
    val_loss, val_acc = validate(...)
    
    # 记录标量
    writer.add_scalar('Loss/train', train_loss, epoch)
    writer.add_scalar('Loss/val', val_loss, epoch)
    writer.add_scalar('Accuracy/train', train_acc, epoch)
    writer.add_scalar('Accuracy/val', val_acc, epoch)
    
    # 记录学习率
    writer.add_scalar('LR', optimizer.param_groups[0]['lr'], epoch)
    
    # 记录模型参数分布
    for name, param in model.named_parameters():
        writer.add_histogram(f'params/{name}', param, epoch)
        if param.grad is not None:
            writer.add_histogram(f'grads/{name}', param.grad, epoch)
 
writer.close()

查看图片

# 记录图片
images, _ = next(iter(train_loader))
img_grid = torchvision.utils.make_grid(images[:8])
writer.add_image('sample_images', img_grid, epoch)
 
# 记录模型图
writer.add_graph(model, images.to(device))

2️⃣ Weights & Biases (WandB) —— 云平台方案

快速开始

pip install wandb
wandb login  # 注册后获取 API Key
import wandb
 
# 初始化
wandb.init(
    project="my-first-project",
    name="exp_lr_1e-3",
    config={
        "learning_rate": 1e-3,
        "batch_size": 64,
        "epochs": 50,
        "model": "ResNet18",
    }
)
 
for epoch in range(wandb.config.epochs):
    train_loss, train_acc = train_one_epoch(...)
    val_loss, val_acc = validate(...)
    
    # 记录指标
    wandb.log({
        "train_loss": train_loss,
        "val_loss": val_loss,
        "train_acc": train_acc,
        "val_acc": val_acc,
        "epoch": epoch,
    })
 
wandb.finish()

WandB 优势

功能说明
☁️ 云端保存所有实验记录云端,随时查看
📊 自动图表Loss/Acc 曲线、参数分布
🔄 实验对比对比不同超参数的效果
📝 自动记录代码、环境、GPU 使用率
🤝 团队协作分享实验链接给同事

3️⃣ 需要监控的指标清单

训练阶段

✅ 训练 Loss     → 应该在稳步下降
✅ 训练 Accuracy  → 应该在逐步上升
✅ 学习率         → 应该按计划衰减
✅ GPU 利用率    → 应该有 70%+(否则 DataLoader 慢)
✅ 显存使用量    → 不能爆显存

验证阶段

✅ 验证 Loss     → 与训练 Loss 差距不应太大
✅ 验证 Accuracy  → 与训练 Accuracy 的差距 = 过拟合程度
✅ 梯度范数      → 太大 => 梯度爆炸,太小 => 梯度消失

4️⃣ 用 tqdm 做终端进度条

from tqdm import tqdm
 
# 简单使用
for i in tqdm(range(100)):
    time.sleep(0.01)
 
# 训练中使用
pbar = tqdm(train_loader, desc=f"Epoch {epoch}")
for batch_x, batch_y in pbar:
    outputs = model(batch_x)
    loss = criterion(outputs, batch_y)
    
    pbar.set_postfix({
        'loss': f'{loss.item():.4f}',
        'lr': f'{optimizer.param_groups[0]["lr"]:.2e}'
    })

5️⃣ 自建简易监控工具

class MetricTracker:
    def __init__(self):
        self.metrics = {}
    
    def update(self, **kwargs):
        for key, value in kwargs.items():
            if key not in self.metrics:
                self.metrics[key] = []
            self.metrics[key].append(value)
    
    def plot(self):
        import matplotlib.pyplot as plt
        fig, axes = plt.subplots(1, 2, figsize=(12, 4))
        
        # Loss
        axes[0].plot(self.metrics.get('train_loss', []), label='train')
        axes[0].plot(self.metrics.get('val_loss', []), label='val')
        axes[0].set_xlabel('Epoch')
        axes[0].set_ylabel('Loss')
        axes[0].legend()
        
        # Accuracy
        axes[1].plot(self.metrics.get('train_acc', []), label='train')
        axes[1].plot(self.metrics.get('val_acc', []), label='val')
        axes[1].set_xlabel('Epoch')
        axes[1].set_ylabel('Accuracy')
        axes[1].legend()
        
        plt.tight_layout()
        plt.show()
    
    def save(self, path='metrics.json'):
        import json
        with open(path, 'w') as f:
            json.dump(self.metrics, f)

📌 一句话总结

TensorBoard 本地够用,WandB 云端更好,tqdm 终端必备


🔗 下一篇

模型评估与验证