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 预处理等）。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！