pytorch训练可视化工具---TensorBoard
一、目的：为什么使用 TensorBoard 调控模型
使用 TensorBoard 可以帮我们：
-
实时查看 loss / acc 曲线 → 判断是否过拟合、欠拟合；
-
对比不同模型或超参数的效果；
-
可视化模型结构 → 帮助调试模型设计；
-
查看权重/梯度分布 → 分析训练稳定性；
-
可视化预测结果、特征图、embedding → 提升模型可解释性；
-
管理实验结果、超参数组合。
二、训练中可视化调控的功能详解（附代码）
1. 可视化 loss / accuracy 曲线
writer.add_scalar("Loss/train", train_loss, epoch)
writer.add_scalar("Loss/val", val_loss, epoch)
writer.add_scalar("Acc/train", train_acc, epoch)
writer.add_scalar("Acc/val", val_acc, epoch)
用途：
-
判断训练过程是否收敛；
-
验证集 loss 高于训练集 → 可能过拟合；
-
loss 一直不下降 → 学习率可能过高或模型设计问题。
2. 可视化模型结构
from torch.utils.tensorboard import SummaryWritermodel = MyModel()
dummy_input = torch.randn(1, 3, 224, 224)
writer.add_graph(model, dummy_input)
用途：
-
检查模型结构是否正确；
-
直观看到各层连接顺序。
3. 可视化权重分布（Histogram）
for name, param in model.named_parameters():writer.add_histogram(name, param, epoch)
用途：
-
观察参数值分布是否爆炸或消失；
-
梯度消失或爆炸时通常 histogram 变化异常。
4. 可视化预测图像、标签、特征图
import torchvision.utils as vutils# 展示输入图像和预测结果
writer.add_images("Input/Image", input_tensor, epoch)
writer.add_images("Predict/Output", output_tensor, epoch)
writer.add_images("GroundTruth/Label", label_tensor, epoch)
用途：
-
视觉任务（如分割、分类）中快速检查模型预测是否合理。
5. 可视化 Embedding（高维向量降维）
# features: [N, D], labels: [N], images: [N, C, H, W]
writer.add_embedding(features, metadata=labels, label_img=images, global_step=epoch)
用途：
-
检查不同类别是否在特征空间中聚类良好；
-
Embedding 层是否学习到有效的表示。
6. 超参数记录与对比（add_hparams）
writer.add_hparams({'lr': 0.001, 'batch_size': 32},{'hparam/accuracy': acc, 'hparam/loss': loss}
)
用途：
-
对比不同超参数下的模型效果；
-
自动生成汇总表格。
7. 可视化学习率变化（学习率调度）
lr = optimizer.param_groups[0]['lr']
writer.add_scalar("LR", lr, epoch)
用途：
-
学习率调度策略是否生效；
-
与 loss 的变化相互印证。
三、训练中综合应用示例（完整代码框架）
from torch.utils.tensorboard import SummaryWriter
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms# 准备
writer = SummaryWriter(log_dir='runs/exp1')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 数据
transform = transforms.ToTensor()
train_loader = torch.utils.data.DataLoader(torchvision.datasets.MNIST('.', train=True, download=True, transform=transform),batch_size=64, shuffle=True)# 模型
model = nn.Sequential(nn.Flatten(),nn.Linear(28*28, 128),nn.ReLU(),nn.Linear(128, 10)
).to(device)# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 添加模型结构
writer.add_graph(model, torch.randn(1, 1, 28, 28).to(device))# 训练
for epoch in range(5):total_loss = 0correct = 0for images, labels in train_loader:images, labels = images.to(device), labels.to(device)outputs = model(images)loss = criterion(outputs, labels)optimizer.zero_grad()loss.backward()optimizer.step()total_loss += loss.item()_, predicted = outputs.max(1)correct += predicted.eq(labels).sum().item()avg_loss = total_loss / len(train_loader)accuracy = correct / len(train_loader.dataset)writer.add_scalar("Loss/train", avg_loss, epoch)writer.add_scalar("Acc/train", accuracy, epoch)writer.add_scalar("LR", optimizer.param_groups[0]['lr'], epoch)# 权重分布for name, param in model.named_parameters():writer.add_histogram(name, param, epoch)# 可视化输入图像img_grid = torchvision.utils.make_grid(images[:16].cpu())writer.add_image("Sample Inputs", img_grid, epoch)writer.close()
四、进阶建议
| 功能 | 说明 |
|---|---|
| 多实验对比 | 使用 SummaryWriter(log_dir=f"runs/lr_{lr}_bs_{bs}") 多次训练 |
| 与 wandb 联动 | 用 wandb 替代 TensorBoard，支持自动超参搜索 |
| TensorBoard.dev | 上传训练记录到云端，便于展示或记录 |
五、总结表格（常用 API）
| 功能 | API |
|---|---|
| 标量值（loss） | add_scalar(tag, value, step) |
| 图像 | add_image(tag, image_tensor, step) |
| 多图像 | add_images(tag, batch_tensor, step) |
| 模型结构 | add_graph(model, input_tensor) |
| 参数直方图 | add_histogram(tag, values, step) |
| 超参对比 | add_hparams(dict, metrics) |
| Embedding | add_embedding(features, labels, images) |