PyTorch + TensorBoard 超实用笔记:从零开始监控你的模型训练
TensorBoard 是深度学习训练中必不可少的“仪表盘”。它可以实时展示损失曲线、准确率变化、模型结构、参数分布甚至图像和音频。 好消息是PyTorch 原生支持 TensorBoard装一个包就能用完全不用写 TensorFlow 代码。本文会带你从零上手 TensorBoard 的核心 API每个功能都配有参数表、代码示例和运行效果说明。这是一篇面向初学者的完整笔记1.SummaryWriter– 创建日志记录器作用创建日志目录TensorBoard 会读取该目录下的所有事件文件。核心参数参数类型默认值说明log_dirstrNone日志保存的文件夹路径。若为None会自动创建runs/下带时间戳的目录。推荐手动指定如logs/exp1。commentstr仅当log_dirNone时有效追加到自动生成的文件夹名后。flush_secsint120每隔多少秒将缓冲区数据强制写入磁盘。训练长时间任务可设为60或30。示例from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(logs/my_experiment) # 最常用 writer.close()启动 TensorBoard在终端中切换到存放logs文件夹的目录执行tensorboard --logdirlogs2.add_scalar– 记录标量作用记录单个数值随步数的变化生成曲线。使用场景训练损失、验证准确率、学习率、自定义指标等。参数参数类型说明tagstr曲线的名字。用/可以分组例如Loss/train和Loss/val会被归入同一组。scalar_valuefloat/int要记录的值y 轴。global_stepint步数x 轴通常用 epoch 或 iteration 编号。walltimefloat可选实际时间戳一般不用。示例记录训练损失和验证准确率for epoch in range(10): train_loss ... # 计算得到 val_acc ... writer.add_scalar(Loss/train, train_loss, epoch) writer.add_scalar(Accuracy/val, val_acc, epoch)add_scalars– 在一张图上画多条曲线需求将训练损失和验证损失放在同一张图中对比。writer SummaryWriter(logs/two_curves) for epoch in range(10): train_loss ... # 计算得到 val_loss ... writer.add_scalars(Loss, {train: train_loss, validation: val_loss}, epoch) writer.close()在 SCALARS 面板中会出现一张名为Loss的图包含两条不同颜色的曲线。3.add_image/add_images– 记录图像作用将图像或一批图像拼成的网格写入 TensorBoard。使用场景检查输入样本、生成模型的输出、特征图、分割掩码等。参数add_image参数类型说明tagstr图像组名称。img_tensortorch.Tensor / numpy.ndarray图像数据。add_image要求形状(C, H, W)或(H, W, C)。global_stepint步数用于观察训练不同阶段的图像变化。dataformatsstr指定维度顺序默认CHW。若张量为(H, W, C)则需设为HWC。常用技巧用torchvision.utils.make_grid将一批图像拼成网格再用add_image记录。示例显示一批样本import torchvision.utils as vutils images, _ next(iter(dataloader)) grid vutils.make_grid(images, nrow8, normalizeTrue) writer.add_image(samples, grid, 0)也使用不同的tag分组展示writer.add_image(cat, cat_img, global_step0, dataformatsHWC) writer.add_image(dog, dog_img, global_step0, dataformatsHWC)4.add_graph– 记录模型计算图作用可视化模型的结构和数据流向。使用场景检查模型定义是否正确调试前向传播形状。参数参数类型说明modeltorch.nn.Module要可视化的模型实例。input_to_modeltorch.Tensor 或 tuple示例输入形状与真实输入一致。TensorBoard 会用它执行一次前向传播来追踪图。verbosebool是否打印详细日志默认False。示例dummy_input torch.randn(1, 1, 28, 28) writer.add_graph(model, dummy_input)打开 GRAPHS 面板可以像看电路图一样展开每个卷积、全连接、残差块检查连接是否正确。5.add_histogram– 记录张量分布直方图作用显示权重、梯度等张量的数值分布随时间的变化。使用场景诊断梯度消失/爆炸观察参数更新是否健康。参数参数类型说明tagstr直方图名称通常用参数名如fc1.weight。valuestorch.Tensor要统计分布的张量任意形状会自动展平。global_stepint步数。binsstr/int分箱方式默认tensorflow。也可指定整数箱数。max_binsint最大箱数当bins为字符串时有效。示例每个 epoch 记录所有参数的权重和梯度for name, param in model.named_parameters(): writer.add_histogram(name, param, epoch) if param.grad is not None: writer.add_histogram(name .grad, param.grad, epoch)怎么看结果HISTOGRAMS 面板正常梯度值集中在 e-2 到 e0 之间分布对称权重直方图逐渐展宽。梯度消失几轮之后梯度直方图缩成一条细线贴在 0 上。梯度爆炸直方图出现很长的尾巴数值远超其他部分。掌握这个你就能很快定位训练失败的根本原因。6.add_hparams– 记录超参数作用系统化地记录一组超参数及其对应的最终指标生成表格和图表。使用场景在你获得该组超参数对应的最终评价指标之后调用通常有几种常见做法单次实验结束后训练完一个模型得到最佳验证集准确率立即调用add_hparams记录这组超参数和最终结果。多次实验的脚本中比如用循环或并行搜索超参数时每次实验跑完就调用一次这样 TensorBoard 的 HParams 插件里会累积多组实验记录。注意不要在每个 epoch 都调用add_hparams是为最终结果设计的。如果你想记录训练过程中指标的变化应该用add_scalar。参数参数类型说明hparam_dictdict超参数字典键为超参数名字符串值为数值或字符串。metric_dictdict指标字典键为指标名建议以hparam/开头值为数值。global_stepint可选步数。示例# 待搜索的超参数组合 learning_rates [0.001, 0.01] batch_sizes [32, 64] epochs 3 # 为了快速演示只训练3轮实际调参可以增加 # 根日志目录 base_log_dir ./runs/mnist_hparam_search for lr in learning_rates: for batch_size in batch_sizes: # 1) 为每组超参数创建独立的子目录 run_name flr_{lr}_bs_{batch_size} log_dir os.path.join(base_log_dir, run_name) # 2) 使用 with 语句自动管理 SummaryWriter with SummaryWriter(log_dir) as writer: print(f\n Running: {run_name} ) # 3) 准备数据加载器不同的 batch_size train_loader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue) val_loader DataLoader(val_dataset, batch_size1000, shuffleFalse) # 4) 创建模型和优化器可加入 dropout 等超参数这里固定 dropout0.2 model SimpleCNN(dropout0.2).to(device) optimizer optim.Adam(model.parameters(), lrlr) # 5) 训练多个 epoch并记录曲线 best_acc 0.0 for epoch in range(1, epochs 1): train(model, device, train_loader, optimizer, epoch, writer) acc validate(model, device, val_loader, writer, epoch) if acc best_acc: best_acc acc # 6) 最后记录超参数和最终验证准确率使用最佳或最终 hparams {lr: lr, batch_size: batch_size} metrics {hparam/val_accuracy: best_acc} writer.add_hparams(hparams, metrics) print(fFinished {run_name}, best accuracy {best_acc:.2f}%)lrbatch_sizeval_accuracy0.001640.930.001320.910.01640.880.01320.85运行脚本后日志会保存在./runs/mnist_hparam_search/下子目录结构为runs/mnist_hparam_search/ ├── lr_0.001_bs_32/ ├── lr_0.001_bs_64/ ├── lr_0.01_bs_32/ └── lr_0.01_bs_64/启动 TensorBoardtensorboard --logdir./runs/mnist_hparam_search浏览器操作打开http://localhost:6006左侧 Runs 列表中你会看到四个可勾选的 Run对应四个子目录。Scalars 面板勾选多个 Run 后train/loss、val/loss、val/accuracy曲线会叠加在同一张图上颜色不同可以清晰对比不同超参数下的收敛速度和最终性能。HParams 面板自动汇总所有 Run 的超参数和最终指标hparam/val_accuracy呈现为一个表格支持排序和交互式图表平行坐标、散点图。8.add_embedding– 高维嵌入可视化深度学习模型尤其是分类任务通常会在最后一层softmax 之前输出一个高维特征向量比如 128 维、512 维。这个向量就是模型对输入数据的内部表示。add_embedding的作用就是把这些高维向量投影到2D 或 3D 空间使用 PCA、t-SNE 或 UMAP然后在 TensorBoard 的PROJECTOR插件中展示出来。你可以用鼠标拖拽旋转/缩放三维空间观察不同类别的样本是否聚集在一起点击某个点查看对应的原始图片或标签搜索特定标签的样本参数writer.add_embedding(mat, metadataNone, label_imgNone, global_stepNone, tagdefault)参数名类型说明matTensororndarray形状为(N, D)的特征矩阵。N是样本数量D是特征维度。必须是浮点数。metadatalistofstr可选。长度为N的标签列表每个元素是该样本的类别名称字符串。如果不提供只显示点坐标。label_imgTensor可选。形状(N, C, H, W)的图像张量。当鼠标悬浮在投影空间的点上时会显示该样本对应的图片缩略图。global_stepint可选。记录当前是第几步/第几个 epoch可用于观察训练过程中特征的可分性如何演变。tagstr可选。嵌入向量的名称默认default。如果你多次调用add_embedding比如不同层或不同时间可以用不同tag区分。注意metadata和label_img可以同时提供也可以只提供其中一个。示例提取测试集特征并可视化features torch.cat(features) # (N, D) labels [fclass_{l} for l in labels.tolist()] writer.add_embedding(features, metadatalabels, label_imgimages, global_step0)打开PROJECTOR面板选择不同的降维算法PCA/t-SNE 等就可以玩起来了。完整ResNet18 CIFAR-10代码import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter import torchvision import torchvision.transforms as transforms import torchvision.utils as vutils # ---------- 模型定义---------- class ResNet18ForCIFAR(nn.Module): def __init__(self, num_classes10): super().__init__() backbone torchvision.models.resnet18(weightstorchvision.models.ResNet18_Weights.IMAGENET1K_V1) backbone.conv1 nn.Conv2d(3, 64, kernel_size3, stride1, padding1, biasFalse) backbone.maxpool nn.Identity() self.features nn.Sequential(*list(backbone.children())[:-1]) self.fc nn.Linear(512, num_classes) def forward(self, x): feat self.features(x) feat feat.view(feat.size(0), -1) out self.fc(feat) return out, feat # 分类结果 特征 # ---------- 数据 ---------- transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_val transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) trainset torchvision.datasets.CIFAR10(./data, trainTrue, downloadTrue, transformtransform_train) valset torchvision.datasets.CIFAR10(./data, trainFalse, downloadTrue, transformtransform_val) train_loader DataLoader(trainset, batch_size128, shuffleTrue, num_workers2) val_loader DataLoader(valset, batch_size128, shuffleFalse, num_workers2) device torch.device(cuda if torch.cuda.is_available() else cpu) model ResNet18ForCIFAR(num_classes10).to(device) # ---------- TensorBoard ---------- writer SummaryWriter(runs/cifar10_demo) # 记录模型图注意dummy input 形状匹配 dummy_input torch.randn(1, 3, 32, 32).to(device) writer.add_graph(model, dummy_input) # 记录一批训练样本图像 sample_images, _ next(iter(train_loader)) grid vutils.make_grid(sample_images[:16], nrow4, normalizeTrue) writer.add_image(train_samples, grid, 0) # ---------- 训练 ---------- optimizer optim.Adam(model.parameters(), lr0.001) criterion nn.CrossEntropyLoss() for epoch in range(5): model.train() train_loss 0.0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs, _ model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() * inputs.size(0) train_loss / len(trainset) writer.add_scalar(Loss/train, train_loss, epoch) for name, param in model.named_parameters(): writer.add_histogram(name, param, epoch) if param.grad is not None: writer.add_histogram(name.grad, param.grad, epoch) print(fEpoch {epoch1}, Loss: {train_loss:.4f}) # ---------- 提取特征并调用 add_embedding ---------- model.eval() feature_list, label_list, image_list [], [], [] with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) _, feats model(inputs) feature_list.append(feats.cpu()) label_list.append(labels.cpu()) image_list.append(inputs.cpu()) features torch.cat(feature_list, dim0) labels torch.cat(label_list, dim0) images torch.cat(image_list, dim0) metadata [fclass_{l} for l in labels.tolist()] embed_writer SummaryWriter(runs/cifar10_embedding) embed_writer.add_embedding(features, metadatametadata, label_imgimages, global_step0) embed_writer.close() writer.close() print(训练完成运行 tensorboard --logdirruns 查看结果)如果本文对你有帮助欢迎点赞 、收藏 ⭐、评论