在 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

参考资源

• https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html
• https://pytorch.org/docs/stable/tensorboard.html