PyTorch学习笔记（五）TensorBoard可视化

本文链接：https://blog.youkuaiyun.com/longrootchen/article/details/105876106

本文介绍了如何在PyTorch中使用TensorBoard进行可视化，包括模型计算图、标量变化、权重直方图和图像数据。通过实例演示了add_graph、add_scalar、add_histogram和add_image等方法的使用，帮助理解训练过程和模型表现。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Environment

OS: macOS Mojave
Python version: 3.7
PyTorch version: 1.4.0
IDE: PyCharm

0. 写在前面

TensorBoard 是 TensorFlow 中常用的可视化工具。从 1.1 版本开始，PyTorch 支持使用 TensorBoard。

安装 TensorBoard

使用 conda install tesnorboard 安装（坑：若 import torch.utils.tensorboard 时提示没有 past 包，则需要安装 future 包（运行 conda install future））

使用 TensorBoard

在 Python 脚本中实例化 torch.utils.tensorboard.SummaryWriter 对象，调用一些实例方法，运行后需要可视化的指标将被记录到硬盘中的 event file
Terminal 中进入到 TensorBoard 创建的 runs 目录的所在路径，即 runs 的上级目录
运行（Terminal） $ tensorboard --logdir=./ 读取 event file 中的内容，在网页端显示
上一步将得到一个地址，从浏览器进入查看

TensorBoard 可视化都是通过 SummaryWriter 类的各种实例方法来实现的，一切从此开始

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(
    log_dir=None,
    comment='',
    purge_step=None,
    max_queue=10,
    flush_secs=120,
    filename_suffix='',
)

1. add_graph 方法

实例方法 add_graph 将模型以计算图的形式进行可视化

writer.add_graph(
    model,  # 模型
    input_to_model=None,  # 输入的数据 batch，随意 torch.randn 一个
    verbose=False
)

2. add_scalar 和 add_scalars 方法

实例方法 add_scalar 和 add_scalars，前者记录单个标量，如可以记录训练损失函数值的变化；后者能够在一幅图中显示多个标量指标，如可以同时记录训练和验证的准确度变化。

writer.add_scalar(
    tag,  # 可视化曲线图的标题
    scalar_value,  # 需要记录的标量值，曲线图的纵坐标
    global_step=None,  # 曲线图的横坐标，通常为对应的迭代数
    walltime=None
)

writer.add_scalars(
    main_tag,  # 可视化曲线图的标题
    tag_scalar_dict,  # 需要可视化的标量的名字和值 {name1: val1, name2: val2, ...}
    global_step=None,
    walltime=None,
)

3. add_histogram 方法

实例方法 add_histogram 绘制直方图，可以用来查看神经网络权重或梯度的分布。

writer.add_histogram(
    tag,
    values,
    global_step=None,
    bins='tensorflow',
    walltime=None,
    max_bins=None,
)

例如可视化一个线性回归模型的训练过程

import os
import torch
from torch.nn import Module, Linear
from torch.nn import MSELoss
from torch.optim import SGD
from torch.utils.tensorboard import SummaryWriter

torch.manual_seed(0)