如何将 TensorBoard 与 PyTorch 结合使用

创建日期：2020 年 4 月 27 日 | 最后更新日期：2024 年 1 月 31 日 | 最后验证日期：2024 年 11 月 5 日

TensorBoard 是一个用于机器学习实验的可视化工具包。TensorBoard 允许跟踪和可视化损失和准确率等指标，可视化模型图，查看直方图，显示图像等等。在本教程中，我们将介绍 TensorBoard 的安装、与 PyTorch 的基本用法，以及如何在 TensorBoard UI 中可视化您记录的数据。

安装

应该安装 PyTorch，以便将模型和指标记录到 TensorBoard 日志目录中。以下命令将通过 Anaconda (推荐) 安装 PyTorch 1.4+

$ conda install pytorch torchvision -c pytorch

或 pip

$ pip install torch torchvision

在 PyTorch 中使用 TensorBoard

现在让我们尝试在 PyTorch 中使用 TensorBoard！在记录任何内容之前，我们需要创建一个 SummaryWriter 实例。

import torch
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()

默认情况下，Writer 将输出到 ./runs/ 目录。

记录标量

在机器学习中，了解损失等关键指标以及它们在训练过程中如何变化非常重要。Scalar 有助于保存每个训练步骤的损失值，或每个 epoch 后的准确率。

要记录标量值，请使用 add_scalar(tag, scalar_value, global_step=None, walltime=None)。例如，让我们创建一个简单的线性回归训练，并使用 add_scalar 记录损失值。

x = torch.arange(-5, 5, 0.1).view(-1, 1)
y = -5 * x + 0.1 * torch.randn(x.size())

model = torch.nn.Linear(1, 1)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr = 0.1)

def train_model(iter):
    for epoch in range(iter):
        y1 = model(x)
        loss = criterion(y1, y)
        writer.add_scalar("Loss/train", loss, epoch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

train_model(10)
writer.flush()

调用 flush() 方法以确保所有待处理事件已写入磁盘。

请参阅 torch.utils.tensorboard 教程以找到更多可以记录的 TensorBoard 可视化类型。

如果您不再需要 summary writer，请调用 close() 方法。

writer.close()

运行 TensorBoard

通过命令行安装 TensorBoard 以可视化您记录的数据

pip install tensorboard

现在，启动 TensorBoard，指定您上面使用的根日志目录。参数 logdir 指向 TensorBoard 将查找可显示的事件文件的目录。TensorBoard 将递归遍历根植于 logdir 的目录结构，查找 .*tfevents.* 文件。

tensorboard --logdir=runs

转到它提供的 URL 或访问 http://localhost:6006/

此仪表板显示了损失和准确率随每个 epoch 的变化情况。您还可以使用它来跟踪训练速度、学习率和其他标量值。比较不同训练运行中的这些指标有助于改进您的模型。

了解更多

• torch.utils.tensorboard 文档

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