注意
跳转至页面底部 下载完整示例代码。
如何将 TensorBoard 与 PyTorch 结合使用#
创建日期：2020年4月27日 | 最后更新：2026年5月9日 | 最后验证：2024年11月5日
TensorBoard 是一个用于机器学习实验的可视化工具包。TensorBoard 允许跟踪和可视化指标（如损失和准确率）、可视化模型图、查看直方图、显示图像等。在本教程中，我们将涵盖 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 变化。您还可以使用它来跟踪训练速度、学习率和其他标量值。比较不同训练运行中的这些指标对于改进模型非常有帮助。