使用 TensorBoard 可视化模型、数据和训练

在 60 Minutes Blitz 中，我们展示了如何加载数据，并把数据送到我们继承 nn.Module 类的模型，在训练数据上训练模型，并在测试集上测试模型。为了看到发生了什么，当模型训练的时候我们打印输出一些统计值获得对模型是否有进展的感觉。我们可以做的比这更好：PyTorch 整合了 TensorBoard，为可视化训练中的神经网络结果的工具。这篇博文说明了它的一些功能，使用可以被 torchvision.datasets 读入 PyTorch 中的 Fashion-MNIST 数据集。

在本文中，我们将学会如何：

以合适的转换（transforms）读入数据（几乎等同于之前的博文）。
设置 TensorBoard
写入 TensorBoard
使用 TensorBoard 检查模型结构
以少量的代码 TensorBoard 创建之前的可视化的交互式版本

尤其在第 5 点，我们将看到：

几个检查我们训练数据的方法
当训练时如何跟踪我们模型的性能（performance）
一旦训练之后我们如何评估我们的模型性能（performance）

我们将开始于相似于训练分类器的模板代码。

# 导包
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
# 转换（transforms）
transform = transforms.Compose(
    [transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5))]
)
# 数据集（datasets）
trainset = torchvision.datasets.FashionMNIST(
    './data', download=True,
    train=True, transform=transform
)
testset = torchvision.datasets.FashionMNIST(
    './data', download=True,
    train=False, transform=transform
)
# 数据加载器（dataloaders）
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=4,
    shuffle=True, num_workers=2
)
testloader = torch.utils.data.DataLoader(
    testset, batch_size=4,
    shuffle=False, num_workers=2
)
# 类别常量（constant for classes）
classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
          'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')
# 显示一张图像的帮助函数
# （用于下面的 `plot_classes_preds` 函数）
def matplotlib_imshow(img, one_channel=False):
    if one_channel:
        img = img.mean(dim=0)
    img = img / 2 + 0.5  # 反规范化（unormalize）
    npimg = img.numpy()
    if one_channel:
        plt.imshow(npimg, cmap='Greys')
    else:
        plt.imshow(np.transpose(npimg, (1, 2, 0)))

如果你下载的比较慢，可是使用我下载好的：本地下载

我们将从这个演示定义一个类似的模型结构，仅做较少的改动解释事实：图片现在是单通道而不是三通道和尺寸为 28x28 而不是 32x32。

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()

我们将定义和之前一样 optimizer 和 criterion。

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

1. 设置 TensorBoard

现在我们将设置 TensorBoard，从 torch.utils 导入 tensorboard 并定义一个 SummaryWriter，我们向 TensorBoard 写入信息的关键对象。

import tensorflow as tf
import tensorboard as tb
tf.io.gfile = tb.compat.tensorflow_stub.io.gfile
from torch.utils.tensorboard import SummaryWriter
# 默认 `log_dir` 是 "runs" —— 在这里我们更具体一些
writer = SummaryWriter("runs/fashion_mnist_experiment_1")

注意上面这一行创建一个 runs/fashion_mnist_experiment_1 文件夹。

2. 写入 TensorBoard

现在，让我们写入一张图片到 TensorBoard —— 确切的来说，一个网格（grid）—— 使用 make_grid。

# 获得一些随机训练图像
dataiter = iter(trainloader)
images, labels = dataiter.next()
# 创建图像网格（grid of images）
img_grid = torchvision.utils.make_grid(images)
# 显示图像
matplotlib_imshow(img_grid, one_channel=True)
# 写入到 TensorBoard
writer.add_image('four_fashion_mnist_images', img_grid)

现在运行。

tensorboard --logdir=runs

从命令行然后导航到 http://localhost:6006/ 应该会看到如下。

现在你知道如何使用 TensorBoard 了！然而，这个例子可被完成在 Jupyter Notebook 中 —— TensorBoard 真正擅长的是创建交互式可视化。我们下一步将介绍这些中的一个，在本博文的最后你会看到几个更多的例子。

3. 使用 TensorBoard 检查模型

TensorBoard 的一个强大的能力就是它可以可视化复杂的模型结构。让我们可视化我们构建的模型。

writer.add_graph(net, images)
writer.close()

现在刷新 TensorBoard 你应该能从上面看到一个“Graphs”就像这样。

继续并双击“Net”以查看它的扩展，查看组成模型的每个单独的操作的详细的视图。

TensorBoard 对于可视化高维数据有着非常方便的特征，比如低维空间中图像数据；我们将在下一步讨论它。

4. 向 TensorBoard 添加“Projector”

我们可以通过 add_embedding 方法可视化更高维度数据的低维表示。

# 辅助函数
def select_n_random(data, labels, n=100):
    '''
    从一个数据集选择 n 个随机数据点和它们对应的标签
    '''
    assert len(data) == len(labels)
    perm = torch.randperm(len(data))
    return data[perm][:n], labels[perm][:n]
# 选择随即图像和它们的目标索引
images, labels = select_n_random(trainset.data, trainset.targets)
# 对每一张图像获得类标签
class_labels = [classes[lab] for lab in labels]
# 载入 embeddings
features = images.view(-1, 28 * 28)
writer.add_embedding(
    features, metadata=class_labels,
    label_img=images.unsqueeze(1)
)
writer.close()

现在，在 TensorBoard 的“Projector”标签里，你可以看到这 100 张图像（每个都是 784 维）投影到三维空间中。除此之外，这还是交互式的：你可以点击并且拖拽旋转三维的投影。最终，几个小提示让可视化更轻松地看到：选择左上方的“color: label”，并且也启用“night mode”，因为它们的背景是白色，这将让图像更轻松地显示出来。

现在我们彻底地检查了我们的数据，现在让我们的展示 TensorBoard 如何更清楚的从训练开始跟踪模型训练和评估。

5. 使用 TensorBoard 跟踪模型训练

在之前的例子中，我们只是每 2000 次迭代就简单地打印输出模型的损失值（loss）。现在，我们将把运行中的损失值（loss）记录到 TensorBoard，以及通过 plot_classes_pred 函数查看模型所做的预测。

# 辅助函数
def images_to_probs(net, images):
    '''
    从一个训练好的网络和一系列图像中生成预测和对应的概率
    '''
    output = net(images)
    # 转换输出概率到预测的类别
    _, preds_tensor = torch.max(output, 1)
    preds = np.squeeze(preds_tensor.numpy())
    return preds, [F.softmax(el, dim=0)[i].item() for i, el in zip(preds, output)]
def plot_classes_preds(net, images, labels):
    '''
    使用一个训练的网络生成 matplotlib 图像，以及从一个数据批量的图片和标签，
    显示出网络的最高的预测和概率以及真实的标签，并且根据预测是否正确信息上色。
    使用 images_to_prbs 函数。
    '''
    preds, probs = images_to_probs(net, images)
    # 画出批量中的图像以及预测的和真实的标签。
    fig = plt.figure(figsize=(12, 6))
    for idx in np.arange(4):
        ax = fig.add_subplot(1, 4, idx + 1, xticks=[], yticks=[])
        matplotlib_imshow(images[idx], one_channel=True)
        ax.set_title("{0}, {1:.1f}%\n(label: {2})".format(
            classes[preds[idx]],
            probs[idx] * 100.0,
            classes[labels[idx]]),
                     color=("green" if preds[idx] == labels[idx].item() else "red")
        )
    return fig

最后，让我们使用与之前的演示相同的代码训练模型，但是将每 1000 批量的结果写入 TensorBoard 而不是打印输出到控制台（console），这一工作使用 add_scalar 函数。

另外，当我们训练时，我们将生成一张图像，显示了包含在数据批量里的四张图像的模型预测和真实结果的对比。

running_loss = 0.0
for epoch in range(1):  # 遍历数据集多次
    for i, data in enumerate(trainloader, 0):
        # 获取输入；data 是一个列表：[inputs, labels]
        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 的数据批量
            # 记录运行的损失值（loss）
            writer.add_scalar('training loss', running_loss / 1000, epoch * len(trainloader) + i)
            # 在随机的数据批量记录一个 Matplotlib 图像，展示模型的预测
            writer.add_figure('predictions vs. actuals',
                             plot_classes_preds(net, inputs, labels),
                             global_step=epoch * len(trainloader) + i)
            running_loss = 0.0
print('Finished Training')

Finished Training

你现在可以在“SCALARS”标签下看到训练的 15,000 次的迭代运行中的 loss。

此外，我们可以看到学习过程中任意的数据批量模型做的预测。在 Images 标签里，滚到 predictions vs. actuals 可视化下方查看；这向我们展示，举个例子，刚过 3000 训练迭代后，即使它还没有训练之后那样自信，模型就有能力区分视觉上有明显区分的类别，比如 shirts，sneakers 和 coats。

在之前的教程里，模型一旦训练完成，我们就看每个类别的准确率；这里，我们将使用 TensorBoard 对每一个类别画出精度—召回率曲线（precision_recall curves），这里有一个很好的解释。

6. 使用 TensorBoard 评估训练好的模型

# 1. 获得预测概率 test_size * num_classes 张量（Tensor）
# 2. 获取 preds 尺寸为 test_size 张量（Tensor）
# 运行时间大约为 10 秒
class_probs = []
class_preds = []
with torch.no_grad():
    for data in testloader:
        images, labels = data
        output = net(images)
        class_probs_batch = [F.softmax(el, dim=0) for el in output]
        _, class_preds_batch = torch.max(output, 1)
        class_probs.append(class_probs_batch)
        class_preds.append(class_preds_batch)
test_probs = torch.cat([torch.stack(batch) for batch in class_probs])
test_preds = torch.cat(class_preds)
# 辅助函数
def add_pr_curve_tensorboard(class_index, test_probs, test_preds, global_step=0):
    '''
    接收从 0 到 9 的 class_index 并且画出对应的 precision-recall curve
    '''
    tensorboard_preds = test_preds == class_index
    tensorboard_probs = test_probs[:, class_index]
    writer.add_pr_curve(classes[class_index], tensorboard_preds, tensorboard_probs, global_step=global_step)
    writer.close()
# 画出所有的 precision-recall curve
for i in range(len(classes)):
    add_pr_curve_tensorboard(i, test_probs, test_preds)

你现在将看到 PR CURVES 菜单标签，其中包含了每一个类别的 precision-recall curves。去四处看看，你将看到一些类别模型有接近 100% 的曲线下的面积（area under the curve），然而其它类别中，这个标记会低一点。

这是一篇关于 TensorBoard 与 PyTorch 的整合的介绍。当然，你可以在 Jupyter Notebook 做 TensorBoard 能做到的事情，但是在 TensorBoard，默认情况下可以得到交互式的视觉效果。