PyTorch 使用 TensorBoard 实现可视化
在 PyTorch 中使用 TensorBoard，可以实现样本数据、模型、训练过程的可视化，能够非常直观地查看在整个训练过程中产生的效果，方便分析和解决一些问题。本文完全根据 PyTorch 官方的 TensorBoard Tutorial 文档进行实践操作，体验 TensorBoard 的基本可视化功能。
首先，需要下载 tensorboard 模块：
pip3 install tensorboard
这样就可以使用 TensorBoard 了，需要通过如下命令，启动 TensorBoard 服务：
tensorboard --logdir=runs
上面 runs 目录是我们写数据的目录，可以根据自己的需要设置目录名称。在使用 PyTorch 过程中需要可视化的话，就把对应的数据写入到这个目录下面的指定文件中，TensorBoard 会直接从该目录下读取并进行可视化。
启动成功后，可以打开浏览器链接 http://localhost:6006/，查看 TensorBoard UI 界面。
我们向 runs 目录中写入数据，直接使用 torch.utils.tensorboard.SummaryWriter 就可以实现，所以先要创建可用的 SummaryWriter 对象，代码如下：
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('runs/fashion_mnist_experiment_2')
这样，就可以在 runs/fashion_mnist_experiment_2 下面为我们进行的一次可视化分析准备数据了，也可以通过指定不同的名称实现多次分析。
数据与模型准备
使用 FashionMNIST 数据集，首先准备 DataSet 和 DataLoader；同时，定义了一个 MyExampleModel 模型类，并设置了 loss function 和优化器，代码如下：
import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
# Define example model
class MyExampleModel(nn.Module):
def __init__(self):
super(MyExampleModel, 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
# Create model, and define loss function and optimizer
model = MyExampleModel()
print(f"My example model: {model}")
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
# Prepare datasets and dataloaders
batch_size = 4
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = torchvision.datasets.FashionMNIST('./data', download=True, train=True, transform=transform)
test_dataset = torchvision.datasets.FashionMNIST('./data', download=True, train=False, transform=transform)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)
# show a batch of example data sample
for X, y in train_dataloader:
print(f"Shape of X [N, C, H, W]: {X.shape}")
print(f"Shape of y: {y.shape} {y.dtype}")
break
后面的一些支持可视化操作收集数据以及实现可视化功能的逻辑，可以基于上面代码中准备的内容来实现。
使用 TensorBoard 可视化模型结构
可视化模型结构，需要将模型 model 作为参数传入，同时也要给模型输入一些数据，这里从 train_dataloader 中迭代一批数据：
# Visualize model structure data_iter = iter(train_dataloader) images, labels = next(data_iter) writer.add_graph(model, images) writer.close()
查看 TensorBoard UI，可以点击 TensorBoard UI 中 GRAPHS 菜单下，看到模型的结构，如下图所示：
点击图中的各个元素，可以查看更详细的信息，非常直观。
使用 TensorBoard 可视化样本数据
筛选一些样本数据点，并写入到 TensorBoard 数据目录中，代码如下：
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]
# select 100 images and their target indices
images, labels = select_n_random(train_dataset.data, train_dataset.targets)
print(f"Length: images = {len(images)}, labels = {len(labels)}")
# constant for classes
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')
# get the class labels for each image
class_labels = [classes[lab] for lab in labels]
print(class_labels)
# log embeddings
features = images.view(-1, 28 * 28)
writer.add_embedding(features, metadata=class_labels, label_img=images.unsqueeze(1))
writer.close()
上面代码，随机从训练集中选了 100 个数据点，然后获取其对应的 label，并输出到 TensorBoard 数据目录中。
查看 TensorBoard UI，可视化结果如图所示：
上图中，将 784 维的高维训练数据，以三维的形式展现，并可以根据 label 进行筛选和分析。
另外，可以根据面板上提供的其它各种选项，进行操作查看。
使用 TensorBoard 可视化模型训练过程
# Train the example model, and collect model metrics
# Train the example model
running_loss = 0.0
print('Start to train model...')
for epoch in range(1): # loop over the dataset multiple times
size = len(train_dataloader.dataset)
for i, (X, y) in enumerate(train_dataloader, 0):
optimizer.zero_grad() # zero the parameter gradients
# forward + backward + optimize
pred = model(X)
loss = loss_fn(pred, y)
loss.backward()
optimizer.step()
loss = loss.item()
running_loss += loss
if i % 1000 == 999:
current = (i + 1) * len(X)
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
# log the running loss
writer.add_scalar('training loss',
running_loss / 1000,
epoch * len(train_dataloader) + i)
running_loss = 0.0
print('Finished Training')
writer.close()
查看 TensorBoard UI，如下图所示，展示了两次训练收集的 loss 数据，并进行可视化趋势跟踪：
可以实时查看训练过程中 loss 变化趋势，从而即时调整模型训练。
使用 TensorBoard 可视化精度与召回率
class_probs = []
class_label = []
with torch.no_grad():
for data in test_dataloader:
images, labels = data
output = model(images)
class_probs_batch = [F.softmax(el, dim=0) for el in output]
class_probs.append(class_probs_batch)
class_label.append(labels)
test_probs = torch.cat([torch.stack(batch) for batch in class_probs])
test_label = torch.cat(class_label)
# helper function
def add_pr_curve_tensorboard(class_index, test_probs, test_label, global_step=0):
# Takes in a "class_index" from 0 to 9 and plots the corresponding precision-recall curve
tensorboard_truth = test_label == class_index
tensorboard_probs = test_probs[:, class_index]
writer.add_pr_curve(classes[class_index],
tensorboard_truth,
tensorboard_probs,
global_step=global_step)
writer.close()
# plot all the pr curves
for i in range(len(classes)):
add_pr_curve_tensorboard(i, test_probs, test_label)
查看 TensorBoard UI，如下图所示：
在 PyTorch 中使用 TensorBoard 其它更多功能，比如：
- 按相关性将图表分组，实现整体分析
- 柱状图
- 散点图
- 嵌入图片
- 网格及 3D 点云图
更详细介绍，可以查看官网文档：https://pytorch.org/docs/stable/tensorboard.html
参考资源