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

原文： https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html

校验者：yearing1017

在 60 分钟闪电战中，我们向您展示了如何加载数据，如何向定义为nn.Module子类的模型提供数据，如何在训练数据上训练该模型，以及在测试数据上对其进行测试。 为了了解发生的情况，我们在模型训练期间打印一些统计信息，以了解训练是否在进行中。 但是，我们可以做得更好：PyTorch 与 TensorBoard 集成，该工具旨在可视化神经网络训练运行的结果。 本教程使用 Fashion-MNIST 数据集说明了其某些功能，可以使用 torchvision.datasets 将其读取到 PyTorch 中。

在本教程中，我们将学习如何：

  1. 读取数据并进行适当的转换(与先前的教程几乎相同）。
  2. 设置 TensorBoard。
  3. 写入 TensorBoard。
  4. 使用 TensorBoard 检查模型架构。
  5. 使用 TensorBoard 来创建我们在上一个教程中创建的可视化的替代版本，代码量更少。

具体来说，在第 5 点，我们将看到：

  • 有两种检查训练数据的方法
  • 在训练模型时如何追踪其性能
  • 在训练完成后如何评估模型的性能。

我们将从 CIFAR-10 教程中类似的样板代码开始：

  1. # imports
  2. import matplotlib.pyplot as plt
  3. import numpy as np
  5. import torch
  6. import torchvision
  7. import torchvision.transforms as transforms
  9. import torch.nn as nn
  10. import torch.nn.functional as F
  11. import torch.optim as optim
  13. # transforms
  14. transform = transforms.Compose(
  15.     [transforms.ToTensor(),
  16.     transforms.Normalize((0.5,), (0.5,))])
  18. # datasets
  19. trainset = torchvision.datasets.FashionMNIST('./data',
  20.     download=True,
  21.     train=True,
  22.     transform=transform)
  23. testset = torchvision.datasets.FashionMNIST('./data',
  24.     download=True,
  25.     train=False,
  26.     transform=transform)
  28. # dataloaders
  29. trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
  30.                                         shuffle=True, num_workers=2)
  32. testloader = torch.utils.data.DataLoader(testset, batch_size=4,
  33.                                         shuffle=False, num_workers=2)
  35. # constant for classes
  36. classes = ('T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
  37.         'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot')
  39. # helper function to show an image
  40. # (used in the `plot_classes_preds` function below)
  41. def matplotlib_imshow(img, one_channel=False):
  42.     if one_channel:
  43.         img = img.mean(dim=0)
  44.     img = img / 2 + 0.5     # unnormalize
  45.     npimg = img.numpy()
  46.     if one_channel:
  47.         plt.imshow(npimg, cmap="Greys")
  48.     else:
  49.         plt.imshow(np.transpose(npimg, (1, 2, 0)))

我们将在该教程中定义一个类似的模型架构，仅需进行少量修改即可说明以下事实：图像现在是一个通道而不是三个通道，是 28x28 而不是 32x32：

  1. class Net(nn.Module):
  2.     def __init__(self):
  3.         super(Net, self).__init__()
  4.         self.conv1 = nn.Conv2d(1, 6, 5)
  5.         self.pool = nn.MaxPool2d(2, 2)
  6.         self.conv2 = nn.Conv2d(6, 16, 5)
  7.         self.fc1 = nn.Linear(16 * 4 * 4, 120)
  8.         self.fc2 = nn.Linear(120, 84)
  9.         self.fc3 = nn.Linear(84, 10)
  11.     def forward(self, x):
  12.         x = self.pool(F.relu(self.conv1(x)))
  13.         x = self.pool(F.relu(self.conv2(x)))
  14.         x = x.view(-1, 16 * 4 * 4)
  15.         x = F.relu(self.fc1(x))
  16.         x = F.relu(self.fc2(x))
  17.         x = self.fc3(x)
  18.         return x
  20. net = Net()

我们将在之前定义相同的optimizer和criterion：

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

1. TensorBoard 设置

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

  1. from torch.utils.tensorboard import SummaryWriter
  3. # default `log_dir` is "runs" - we'll be more specific here
  4. writer = SummaryWriter('runs/fashion_mnist_experiment_1')

请注意，仅此行会创建一个runs/fashion_mnist_experiment_1文件夹。

2.写入 TensorBoard

现在，使用 make_grid 将图像写入到 TensorBoard 中，具体来说就是网格。

  1. # get some random training images
  2. dataiter = iter(trainloader)
  3. images, labels = dataiter.next()
  5. # create grid of images
  6. img_grid = torchvision.utils.make_grid(images)
  8. # show images
  9. matplotlib_imshow(img_grid, one_channel=True)
  11. # write to tensorboard
  12. writer.add_image('four_fashion_mnist_images', img_grid)

现在运行

  1. tensorboard --logdir=runs

从命令行，然后导航到 https：// localhost：6006 应该显示以下内容。

intermediate/../../_static/img/tensorboard_first_view.png

现在您知道如何使用 TensorBoard 了！ 但是，此示例可以在 Jupyter Notebook 中完成-TensorBoard 真正擅长的地方是创建交互式可视化。 接下来，我们将介绍其中之一，并在本教程结束时介绍更多内容。

3.使用 TensorBoard 检查模型

TensorBoard 的优势之一是其可视化复杂模型结构的能力。 让我们可视化我们构建的模型。

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

现在刷新 TensorBoard 后，您应该会看到一个“ Graphs”标签，如下所示：

intermediate/../../_static/img/tensorboard_model_viz.png

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

TensorBoard 具有非常方便的功能，可在低维空间中可视化高维数据，例如图像数据； 接下来我们将介绍。

4.在 TensorBoard 中添加一个“投影仪”

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

  1. # helper function
  2. def select_n_random(data, labels, n=100):
  3.     '''
  4.     Selects n random datapoints and their corresponding labels from a dataset
  5.     '''
  6.     assert len(data) == len(labels)
  8.     perm = torch.randperm(len(data))
  9.     return data[perm][:n], labels[perm][:n]
  11. # select random images and their target indices
  12. images, labels = select_n_random(trainset.data, trainset.targets)
  14. # get the class labels for each image
  15. class_labels = [classes[lab] for lab in labels]
  17. # log embeddings
  18. features = images.view(-1, 28 * 28)
  19. writer.add_embedding(features,
  20.                     metadata=class_labels,
  21.                     label_img=images.unsqueeze(1))
  22. writer.close()

现在，在 TensorBoard 的“投影仪”选项卡中，您可以看到这 100 张图像-每个图像 784 维-向下投影到三维空间中。 此外，这是交互式的：您可以单击并拖动以旋转三维投影。 最后，有两个技巧可以使可视化效果更容易看到：在左上方选择“颜色：标签”，并启用“夜间模式”，这将使图像更容易看到，因为它们的背景是白色的：

intermediate/../../_static/img/tensorboard_projector.png

现在我们已经彻底检查了我们的数据，接下来让我们展示TensorBoard 如何从训练开始就可以使跟踪模型的训练和评估更加清晰。

5.使用 TensorBoard 跟踪模型训练

在前面的示例中，我们仅每 2000 次迭代打印该模型的运行损失。 现在，我们将运行损失记录到 TensorBoard 中，并通过plot_classes_preds函数查看模型所做的预测。

  1. # helper functions
  3. def images_to_probs(net, images):
  4.     '''
  5.     Generates predictions and corresponding probabilities from a trained
  6.     network and a list of images
  7.     '''
  8.     output = net(images)
  9.     # convert output probabilities to predicted class
  10.     _, preds_tensor = torch.max(output, 1)
  11.     preds = np.squeeze(preds_tensor.numpy())
  12.     return preds, [F.softmax(el, dim=0)[i].item() for i, el in zip(preds, output)]
  14. def plot_classes_preds(net, images, labels):
  15.     '''
  16.     Generates matplotlib Figure using a trained network, along with images
  17.     and labels from a batch, that shows the network's top prediction along
  18.     with its probability, alongside the actual label, coloring this
  19.     information based on whether the prediction was correct or not.
  20.     Uses the "images_to_probs" function.
  21.     '''
  22.     preds, probs = images_to_probs(net, images)
  23.     # plot the images in the batch, along with predicted and true labels
  24.     fig = plt.figure(figsize=(12, 48))
  25.     for idx in np.arange(4):
  26.         ax = fig.add_subplot(1, 4, idx+1, xticks=[], yticks=[])
  27.         matplotlib_imshow(images[idx], one_channel=True)
  28.         ax.set_title("{0}, {1:.1f}%\n(label: {2})".format(
  29.             classes[preds[idx]],
  30.             probs[idx] * 100.0,
  31.             classes[labels[idx]]),
  32.                     color=("green" if preds[idx]==labels[idx].item() else "red"))
  33.     return fig

最后，让我们使用与之前教程相同的模型训练代码来训练模型，但是每 1000 批将结果写入 TensorBoard，而不是打印到控制台。 这是通过 add_scalar 函数完成的。

此外，在训练过程中，我们将生成一幅图像，显示该批次中包含的四幅图像的模型预测与实际结果。

  1. running_loss = 0.0
  2. for epoch in range(1):  # loop over the dataset multiple times
  4.     for i, data in enumerate(trainloader, 0):
  6.         # get the inputs; data is a list of [inputs, labels]
  7.         inputs, labels = data
  9.         # zero the parameter gradients
  10.         optimizer.zero_grad()
  12.         # forward + backward + optimize
  13.         outputs = net(inputs)
  14.         loss = criterion(outputs, labels)
  15.         loss.backward()
  16.         optimizer.step()
  18.         running_loss += loss.item()
  19.         if i % 1000 == 999:    # every 1000 mini-batches...
  21.             # ...log the running loss
  22.             writer.add_scalar('training loss',
  23.                             running_loss / 1000,
  24.                             epoch * len(trainloader) + i)
  26.             # ...log a Matplotlib Figure showing the model's predictions on a
  27.             # random mini-batch
  28.             writer.add_figure('predictions vs. actuals',
  29.                             plot_classes_preds(net, inputs, labels),
  30.                             global_step=epoch * len(trainloader) + i)
  31.             running_loss = 0.0
  32. print('Finished Training')

现在，您可以查看“标量”选项卡，以查看在 15,000 次训练迭代中绘制的运行损失：

intermediate/../../_static/img/tensorboard_scalar_runs.png

此外，我们可以查看整个学习过程中模型在任意批次上所做的预测。 查看“图像”选项卡，然后在“预测与实际”可视化条件下向下滚动以查看此内容； 这向我们表明，例如，仅经过 3000 次训练迭代，该模型就能够区分出视觉上截然不同的类，例如衬衫，运动鞋和外套，尽管它并没有像后来的训练那样有信心：

intermediate/../../_static/img/tensorboard_images.png

在之前的教程中，我们研究了模型训练后的每类准确性； 在这里，我们将使用 TensorBoard 绘制每个类的精度回召曲线(很好的解释在这里）。

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

  1. # 1\. gets the probability predictions in a test_size x num_classes Tensor
  2. # 2\. gets the preds in a test_size Tensor
  3. # takes ~10 seconds to run
  4. class_probs = []
  5. class_preds = []
  6. with torch.no_grad():
  7.     for data in testloader:
  8.         images, labels = data
  9.         output = net(images)
  10.         class_probs_batch = [F.softmax(el, dim=0) for el in output]
  11.         _, class_preds_batch = torch.max(output, 1)
  13.         class_probs.append(class_probs_batch)
  14.         class_preds.append(class_preds_batch)
  16. test_probs = torch.cat([torch.stack(batch) for batch in class_probs])
  17. test_preds = torch.cat(class_preds)
  19. # helper function
  20. def add_pr_curve_tensorboard(class_index, test_probs, test_preds, global_step=0):
  21.     '''
  22.     Takes in a "class_index" from 0 to 9 and plots the corresponding
  23.     precision-recall curve
  24.     '''
  25.     tensorboard_preds = test_preds == class_index
  26.     tensorboard_probs = test_probs[:, class_index]
  28.     writer.add_pr_curve(classes[class_index],
  29.                         tensorboard_preds,
  30.                         tensorboard_probs,
  31.                         global_step=global_step)
  32.     writer.close()
  34. # plot all the pr curves
  35. for i in range(len(classes)):
  36.     add_pr_curve_tensorboard(i, test_probs, test_preds)

现在，您将看到一个“ PR Curves”选项卡，其中包含每个类别的精度回召曲线。 继续戳一下； 您会发现在某些类别中，模型的“曲线下面积”接近 100％，而在另一些类别中，该面积更低：

intermediate/../../_static/img/tensorboard_pr_curves.png

这是 TensorBoard 和 PyTorch 与之集成的介绍。 当然，您可以在 Jupyter Notebook 中完成 TensorBoard 所做的所有操作，但是使用 TensorBoard，默认情况下，您还可以获得交互式视觉效果。