注意
转到末尾 下载完整示例代码。
如何在 PyTorch 中使用 TensorBoard#
创建于：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/ 目录。
记录标量#
在机器学习中，了解关键指标（例如损失）及其在训练期间的变化非常重要。标量有助于保存每个训练步骤的损失值，或每个 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 可视化类型。
如果您不再需要摘要写入器，请调用 close() 方法。
writer.close()
运行 TensorBoard#
通过命令行安装 TensorBoard 以可视化您记录的数据
pip install tensorboard
现在，启动 TensorBoard，指定您上面使用的根日志目录。logdir 参数指向 TensorBoard 将查找事件文件以供显示所在的目录。TensorBoard 将递归遍历以 logdir 为根的目录结构，查找 .*tfevents.* 文件。
tensorboard --logdir=runs
转到它提供的 URL 或 https://:6006/
此仪表板显示了损失和准确性如何随每个 epoch 变化。您还可以使用它来跟踪训练速度、学习率和其他标量值。比较不同训练运行中的这些指标有助于改进您的模型。
