注意
跳转至页面底部下载完整示例代码。
如何将 TensorBoard 与 PyTorch 配合使用#
创建日期：2020年4月27日 | 最后更新：2026年5月9日 | 最后验证：2024年11月5日
TensorBoard 是一个用于机器学习实验的可视化工具包。TensorBoard 允许跟踪和可视化损失（loss）和准确率（accuracy）等指标、可视化模型图、查看直方图、显示图像等等。在本教程中，我们将介绍 TensorBoard 的安装、在 PyTorch 中的基本用法，以及如何可视化你在 TensorBoard UI 中记录的数据。
安装#
需要安装 PyTorch 才能将模型和指标记录到 TensorBoard 日志目录中。通过以下命令安装 PyTorch：
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 或 https://:6006/
此仪表板显示了损失和准确率随每个 epoch 的变化情况。你还可以使用它来跟踪训练速度、学习率和其他标量值。比较不同训练运行中的这些指标有助于改进你的模型。