当前位置：首页 > news >正文

pytorch小记（三十三）：PyTorch 使用 TensorBoard 可视化训练过程（含完整示例）

news 2025/8/1 17:48:55

pytorch小记（三十三）：深PyTorch 使用 TensorBoard 可视化训练过程（含完整示例）

PyTorch 使用 TensorBoard 可视化训练过程（含完整示例）
- 📌 一、什么是 TensorBoard？
- ✅ 二、SummaryWriter 是什么？
- 🛠 三、常见用法示例
- - 1. 初始化 writer
  - 2. 写入标量（scalar）
  - 3. 写入图像
  - 4. 写入模型结构图
  - 5. 写入参数直方图（查看收敛趋势）
  - 6. 关闭 writer（必须）
- 🚀 四、完整 PyTorch 示例（MNIST 分类）
- 📺 五、运行 TensorBoard
- 🧠 六、常用方法速查表
- ✅ 七、总结

PyTorch 使用 TensorBoard 可视化训练过程（含完整示例）

TensorBoard 是深度学习训练中最重要的可视化工具之一，本文将详细介绍如何在 PyTorch 中使用 SummaryWriter 写入训练过程数据，并通过 TensorBoard 进行可视化，附完整训练示例。

📌 一、什么是 TensorBoard？

TensorBoard 是一个可视化工具，最初由 TensorFlow 团队开发，现在也被 PyTorch 原生支持。它可以帮助我们直观地监控训练过程，比如：

Loss、Accuracy 的变化趋势
模型结构图（Computation Graph）
权重分布直方图（Histogram）
图像输入和中间特征图
甚至音频、视频等数据

在 PyTorch 中，我们只需使用：

from torch.utils.tensorboard import SummaryWriter

即可引入官方支持的可视化工具。

✅ 二、SummaryWriter 是什么？

SummaryWriter 是 PyTorch 提供的一个类，用于将数据写入磁盘，供 TensorBoard 读取。你可以用它记录：

scalar（标量，如 loss, acc）
image（图像）
graph（模型结构）
histogram（参数分布）

🛠 三、常见用法示例

1. 初始化 writer

from torch.utils.tensorboard import SummaryWriterwriter = SummaryWriter(log_dir='runs/exp1')  # 会自动创建 runs/exp1 目录

2. 写入标量（scalar）

writer.add_scalar("Loss/train", loss_value, epoch)
writer.add_scalar("Accuracy/val", acc_val, epoch)

3. 写入图像

writer.add_images("Input Images", images, global_step=0)

4. 写入模型结构图

writer.add_graph(model, input_tensor)

5. 写入参数直方图（查看收敛趋势）

for name, param in model.named_parameters():writer.add_histogram(name, param, epoch)

6. 关闭 writer（必须）

writer.close()

🚀 四、完整 PyTorch 示例（MNIST 分类）

以下是一个完整的训练+TensorBoard可视化代码：

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter# 数据加载
transform = transforms.ToTensor()
trainset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)# 模型定义
model = nn.Sequential(nn.Flatten(),nn.Linear(28*28, 128),nn.ReLU(),nn.Linear(128, 10)
)criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 初始化 TensorBoard Writer
writer = SummaryWriter('runs/mnist_demo')# 训练过程
for epoch in range(5):running_loss = 0.0for i, (inputs, labels) in enumerate(trainloader):optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()avg_loss = running_loss / len(trainloader)writer.add_scalar('Loss/train', avg_loss, epoch)print(f"Epoch [{epoch+1}/5], Loss: {avg_loss:.4f}")# 写入模型结构
sample_input = torch.rand(1, 1, 28, 28)
writer.add_graph(model, sample_input)writer.close()