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

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

在 60 分钟突击中，我们向您展示了如何加载数据，如何通过定义为nn.Module子类的模型提供数据，如何在训练数据上训练该模型以及在测试数据上对其进行测试。 为了了解发生的情况，我们在模型训练期间打印一些统计数据，以了解训练是否在进行中。 但是，我们可以做得更好：PyTorch 与 TensorBoard 集成在一起，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 笔记本中完成-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 笔记本中完成 TensorBoard 的所有操作，但是使用 TensorBoard 时，默认情况下会获得交互式的视觉效果。