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

墨绿色的摆渡人

733人浏览 · 2025-07-30 18:08:06

墨绿色的摆渡人 · 2025-07-30 18:08:06 发布

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

PyTorch 使用 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 SummaryWriter

writer = 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.0
    for 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()