PyTorch TensorBoard 支持
你可以观看下方视频（或前往 YouTube 观看）跟随教程操作。
准备工作
运行本教程前，你需要安装 PyTorch、TorchVision、Matplotlib 和 TensorBoard。
使用 conda 安装：
bash
conda install pytorch torchvision -c pytorch
conda install matplotlib tensorboard使用 pip 安装：
bash
pip install torch torchvision matplotlib tensorboard依赖安装完成后，在安装这些库的 Python 环境中重启本 notebook。
简介
在本 notebook 中，我们将基于 Fashion-MNIST 数据集训练一个改进版的 LeNet-5 模型。Fashion-MNIST 是一组包含各类服装的图像切片数据集，包含 10 个类别标签，用于标识图像中服装的类型。
python
# PyTorch 模型和训练必备库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
# 图像数据集和图像处理库
import torchvision
import torchvision.transforms as transforms
# 图像展示库
import matplotlib.pyplot as plt
import numpy as np
# PyTorch TensorBoard 支持库
from torch.utils.tensorboard import SummaryWriter
# 如果你使用的环境中安装了 TensorFlow（例如 Google Colab），
# 取消以下代码的注释以避免将嵌入保存到 TensorBoard 目录时出现的 bug
# import tensorflow as tf
# import tensorboard as tb
# tf.io.gfile = tb.compat.tensorflow_stub.io.gfile在 TensorBoard 中展示图像
首先，我们将从数据集中提取样本图像并添加到 TensorBoard：
python
# 加载数据集并预处理
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))])
# 将训练集和验证集分别存储在 ./data 目录下
training_set = torchvision.datasets.FashionMNIST('./data',
download=True,
train=True,
transform=transform)
validation_set = torchvision.datasets.FashionMNIST('./data',
download=True,
train=False,
transform=transform)
training_loader = torch.utils.data.DataLoader(training_set,
batch_size=4,
shuffle=True,
num_workers=2)
validation_loader = torch.utils.data.DataLoader(validation_set,
batch_size=4,
shuffle=False,
num_workers=2)
# 类别标签
classes = ('T恤/上衣', '裤子', '套头衫', '连衣裙', '外套',
'凉鞋', '衬衫', '运动鞋', '包', '短靴')
# 图像可视化辅助函数
def matplotlib_imshow(img, one_channel=False):
if one_channel:
img = img.mean(dim=0)
img = img / 2 + 0.5 # 反归一化
npimg = img.numpy()
if one_channel:
plt.imshow(npimg, cmap="Greys")
else:
plt.imshow(np.transpose(npimg, (1, 2, 0)))
# 提取一个批次（4张）的图像
dataiter = iter(training_loader)
images, labels = next(dataiter)
# 创建图像网格并展示
img_grid = torchvision.utils.make_grid(images)
matplotlib_imshow(img_grid, one_channel=True)
shell
tensorboardyt tutorial
0%| | 0.00/26.4M [00:00<?, ?B/s]
0%| | 65.5k/26.4M [00:00<01:12, 364kB/s]
1%| | 197k/26.4M [00:00<00:45, 576kB/s]
3%|▎ | 819k/26.4M [00:00<00:13, 1.88MB/s]
12%|█▏ | 3.28M/26.4M [00:00<00:03, 6.51MB/s]
35%|███▌ | 9.27M/26.4M [00:00<00:01, 16.1MB/s]
58%|█████▊ | 15.4M/26.4M [00:01<00:00, 22.2MB/s]
80%|████████ | 21.3M/26.4M [00:01<00:00, 29.9MB/s]
93%|█████████▎| 24.7M/26.4M [00:01<00:00, 26.2MB/s]
100%|██████████| 26.4M/26.4M [00:01<00:00, 19.4MB/s]
0%| | 0.00/29.5k [00:00<?, ?B/s]
100%|██████████| 29.5k/29.5k [00:00<00:00, 326kB/s]
0%| | 0.00/4.42M [00:00<?, ?B/s]
1%|▏ | 65.5k/4.42M [00:00<00:12, 361kB/s]
5%|▌ | 229k/4.42M [00:00<00:06, 680kB/s]
20%|██ | 885k/4.42M [00:00<00:01, 2.02MB/s]
81%|████████ | 3.57M/4.42M [00:00<00:00, 7.05MB/s]
100%|██████████| 4.42M/4.42M [00:00<00:00, 6.07MB/s]
0%| | 0.00/5.15k [00:00<?, ?B/s]
100%|██████████| 5.15k/5.15k [00:00<00:00, 58.4MB/s]上文我们使用 TorchVision 和 Matplotlib 创建了输入数据批次的可视化网格。接下来，我们使用 SummaryWriter 的 add_image() 方法将图像记录到 TensorBoard 日志中，并调用 flush() 确保数据立即写入磁盘：
python
# 默认 log_dir 参数为 "runs" - 但最好显式指定
# torch.utils.tensorboard.SummaryWriter 已在上方导入
writer = SummaryWriter('runs/fashion_mnist_experiment_1')
# 将图像数据写入 TensorBoard 日志目录
writer.add_image('四张 Fashion-MNIST 图像', img_grid)
writer.flush()
# 如需查看，在命令行启动 TensorBoard：
# tensorboard --logdir=runs
# ...然后在浏览器中打开 http://localhost:6006/如果你在命令行启动 TensorBoard 并在新的浏览器标签页中打开（通常地址为 localhost:6006），你应该能在 IMAGES 标签下看到这个图像网格。
绘制标量可视化训练过程
TensorBoard 可用于跟踪训练进度和效果。接下来，我们将运行训练循环、跟踪一些指标，并将数据保存供 TensorBoard 使用。
首先定义用于分类图像的模型，以及训练所需的优化器和损失函数：
python
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 4 * 4, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 4 * 4)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)现在我们训练一个 epoch，并每 1000 个批次评估一次训练集和验证集的损失：
python
print(len(validation_loader))
for epoch in range(1): # 多轮遍历数据集
running_loss = 0.0
for i, data in enumerate(training_loader, 0):
# 基础训练循环
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 1000 == 999: # 每 1000 个小批次...
print('批次 {}'.format(i + 1))
# 在验证集上验证
running_vloss = 0.0
# 评估模式下，模型的某些特定操作会被省略（例如 dropout 层）
net.train(False) # 切换到评估模式（如关闭正则化）
for j, vdata in enumerate(validation_loader, 0):
vinputs, vlabels = vdata
voutputs = net(vinputs)
vloss = criterion(voutputs, vlabels)
running_vloss += vloss.item()
net.train(True) # 切换回训练模式（如开启正则化）
avg_loss = running_loss / 1000
avg_vloss = running_vloss / len(validation_loader)
# 记录每个批次的平均运行损失
writer.add_scalars('训练损失 vs 验证损失',
{ '训练' : avg_loss, '验证' : avg_vloss },
epoch * len(training_loader) + i)
running_loss = 0.0
print('训练完成')
writer.flush()
yaml
2500
批次 1000
批次 2000
批次 3000
批次 4000
批次 5000
批次 6000
批次 7000
批次 8000
批次 9000
批次 10000
批次 11000
批次 12000
批次 13000
批次 14000
批次 15000
训练完成切换到你打开的 TensorBoard 页面，查看 SCALARS 标签页。
可视化模型结构
TensorBoard 也可用于查看模型中的数据流向。为此，调用 add_graph() 方法并传入模型和样本输入：
python
# 再次获取一个小批次的图像
dataiter = iter(training_loader)
images, labels = next(dataiter)
# add_graph() 会跟踪样本输入在模型中的流向，
# 并将其渲染为图形
writer.add_graph(net, images)
writer.flush()切换到 TensorBoard 后，你应该能看到 GRAPHS 标签页。双击 "NET" 节点可以查看模型内部的层和数据流向。
使用嵌入可视化数据集
我们使用的 28×28 图像切片可以建模为 784 维向量（28 * 28 = 784）。将其投影到低维表示中会很有启发性。add_embedding() 方法会将一组数据投影到方差最大的三个维度上，并将其显示为交互式 3D 图表。add_embedding() 方法会自动完成此操作，投影到方差最大的三个维度。
接下来，我们选取一部分数据并生成这样的嵌入：
python
# 随机选择 n 个数据和对应的标签
def select_n_random(data, labels, n=100):
assert len(data) == len(labels)
perm = torch.randperm(len(data))
return data[perm][:n], labels[perm][:n]
# 提取随机子集数据
images, labels = select_n_random(training_set.data, training_set.targets)
# 获取每张图像的类别标签
class_labels = [classes[label] for label in labels]
# 记录嵌入
features = images.view(-1, 28 * 28)
writer.add_embedding(features,
metadata=class_labels,
label_img=images.unsqueeze(1))
writer.flush()
writer.close()现在切换到 TensorBoard 并选择 PROJECTOR 标签页，你应该能看到投影的 3D 表示。你可以旋转和缩放模型，从不同尺度查看，并观察投影数据中的模式和标签聚类情况。
为了获得更好的可视效果，建议：
- 在左侧的 "Color by" 下拉菜单中选择 "label"。
- 点击顶部的夜间模式图标，将浅色图像显示在深色背景上。
其他资源
更多信息请参考：
- PyTorch 官方文档：torch.utils.tensorboard.SummaryWriter
- PyTorch.org 教程中的 TensorBoard 教程内容
- TensorBoard 官方文档
脚本总运行时间：(1 分 50.462 秒)
总结
- 本教程完整演示了PyTorch结合TensorBoard的核心用法，包括图像可视化、训练损失跟踪、模型结构可视化和数据集嵌入可视化；
- 关键操作包括SummaryWriter的初始化、add_image/add_scalars/add_graph/add_embedding等核心API的使用；
- 所有代码可直接运行，需注意安装依赖库并通过
tensorboard --logdir=runs启动TensorBoard查看结果。