Ubuntu上PyTorch的可视化工具如何使用 - 问答

Ubuntu上PyTorch可视化工具使用指南

一工具选型与适用场景

TensorBoard：记录并可视化训练过程中的标量（如Loss/Accuracy）、图像、直方图、计算图等；适合长期跟踪实验与对比多个运行。
Matplotlib/Seaborn/Pandas：训练后在本地绘制曲线、分布、热力图等，适合快速分析与报告输出。
Visdom：Facebook开源，强调实时可视化，适合边训练边观察曲线、图像等。
Torchinfo：模型结构与参数量概览（如输入形状、输出形状、参数量）。
PyTorchviz：可视化计算图（依赖系统级 Graphviz）。
Netron：可视化模型结构（支持多种框架与格式）。

二快速上手 TensorBoard

安装与导入
- 建议安装与当前 PyTorch 版本匹配的 TensorBoard（现代版本自带或通过 pip 安装）：
  - pip 安装：pip install tensorboard
  - 旧版（如 PyTorch ≤1.6）可用 tensorboardX：pip install tensorboardX
- 导入：
  - 新版：from torch.utils.tensorboard import SummaryWriter
  - 旧版：from tensorboardX import SummaryWriter

记录标量与图像

from torch.utils.tensorboard import SummaryWriter
import torch, torchvision

writer = SummaryWriter(log_dir="runs/exp1")

for epoch in range(10):
    loss = torch.rand(1).item()
    acc  = torch.rand(1).item()
    writer.add_scalar("Loss/train", loss, epoch)
    writer.add_scalar("Acc/train",  acc,  epoch)

    # 记录示例图像（C×H×W，范围0~1）
    img = torch.rand(3, 32, 32)
    writer.add_image("sample/input", img, epoch)

writer.close()

启动与查看
- 终端启动：tensorboard --logdir=runs
- 浏览器访问：http://localhost:6006
小提示
- 多实验用不同子目录（如 runs/exp1、runs/exp2），便于对比。
- 若端口被占用，可用：tensorboard --logdir=runs --port=6007。

三实时曲线与图像 Visdom

安装与启动
- 安装：pip install visdom
- 启动服务：python -m visdom.server
- 浏览器访问：http://localhost:8097

绘制曲线与图像

import numpy as np
from visdom import Visdom

viz = Visdom(env="main")

# 实时曲线
x = np.linspace(0, 4*np.pi, 200)
y = np.sin(x)
viz.line(Y=y, X=x, win="sin", opts=dict(title="sin(x)"))

# 图像（H×W×C -> C×H×W）
img = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
viz.image(img.transpose(2, 0, 1), win="rand_img", opts=dict(title="Random Image"))

适用场景
- 训练过程中实时观察 Loss/Accuracy 曲线、特征可视化、样本展示等。

四模型结构与计算图可视化

Torchinfo（模型概览）

from torchinfo import summary
import torchvision.models as models

model = models.resnet18()
summary(model, input_size=(1, 3, 224, 224))

PyTorchviz（计算图，需系统 Graphviz）

安装系统依赖（Ubuntu）：sudo apt-get install graphviz
安装 Python 包：pip install torchviz

示例代码：

import torch
from torchviz import make_dot

x = torch.randn(1, 3, 224, 224, requires_grad=True)
model = models.resnet18()
y = model(x)
loss = y.mean()
loss.backward()

dot = make_dot(y, params=dict(model.named_parameters()))
dot.render("resnet18_graph", format="png", cleanup=True)  # 生成 resnet18_graph.png

Netron（模型结构浏览器）
- 安装：pip install netron
- 使用：
  - 代码启动：netron.start(‘model.pth’)
  - 或命令行：netron model.pth（默认端口 8080，浏览器访问 http://localhost:8080）。

五训练监控与资源查看

GPU 监控
- 实时查看：watch -n 1 nvidia-smi（每秒刷新一次），关注显存占用、利用率、温度等。
CPU/内存监控
- 交互式：sudo apt install htop && htop
- 基础：top

在训练脚本内监控

import psutil, os

pid = os.getpid()
proc = psutil.Process(pid)
print(f"CPU Usage: {proc.cpu_percent(interval=1.0)}%")
print(f"Memory RSS: {proc.memory_info().rss / 1024**2:.2f} MB")

与 TensorBoard 结合
- 将关键指标写入 SummaryWriter，配合 nvidia-smi/htop 定位瓶颈（数据加载、GPU 计算、CPU 预处理等）。

0 赞

0 踩