Linux上PyTorch可视化工具如何使用

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

一、常用工具与适用场景

• TensorBoard / torch.utils.tensorboard.SummaryWriter：记录并可视化训练过程中的标量（如loss/accuracy）、图像、直方图、计算图等，适合系统化对比实验与排查训练问题。
• Visdom：Facebook开源的实时可视化工具，支持折线、散点、热图、图像、文本等，适合边训练边看曲线的交互式场景。
• Netron：模型结构可视化，支持PyTorch等多种框架导出的模型文件，便于快速查看网络层与形状。
• torchinfo：结构化打印与汇总模型信息（层数、参数量、每层输出形状、FLOPs估算等），便于模型复杂度评估。
• PyTorchViz：生成模型计算图（依赖Graphviz），便于理解数据流向与模块依赖。
• Matplotlib / Seaborn / Pandas：训练日志的本地分析与可视化，适合轻量绘图与报表。

二、TensorBoard快速上手

• 安装与启动
  • 安装：pip install tensorboard（PyTorch自带torch.utils.tensorboard，无需额外安装tensorboardx）。
  • 启动：tensorboard --logdir=runs --port=6006（默认端口6006，可更换端口与host）。
• 记录训练指标与图像
  • 代码示例：

    from torch.utils.tensorboard import SummaryWriterimport torch, torchvisionwriter = SummaryWriter(log_dir="runs/exp1")# 模型与数据model = torchvision.models.resnet18(False).cuda()x = torch.randn(4, 3, 224, 224).cuda()# 标量：损失与准确率for epoch in range(5):loss = torch.rand(1).cuda()acc = torch.rand(1).cuda()writer.add_scalar("Loss/train", loss, epoch)writer.add_scalar("Acc/train", acc, epoch)# 直方图：参数分布for name, param in model.named_parameters():writer.add_histogram(f"param/{name}", param, epoch)# 图像：输入样例网格if epoch == 0:grid = torchvision.utils.make_grid(x, nrow=2)writer.add_image("input", grid, epoch)writer.close()

  • 在浏览器访问：http://localhost:6006 查看各项面板。

三、Visdom实时可视化

• 安装与启动
  • 安装：pip install visdom
  • 启动：python -m visdom.server -port 8097（默认端口8097）。
• 记录训练曲线、图像与文本
  • 代码示例：

    import visdom, numpy as np, timeviz = visdom.Visdom(env="demo")# 单条曲线for step in range(100):loss = np.random.randn()viz.line([loss], [step], win="loss", update="append", opts=dict(title="Train Loss"))# 多条曲线：loss与accfor step in range(100):loss = abs(np.random.randn()) + 1acc= abs(np.random.randn())viz.line([[loss, acc]], [step], win="metrics", update="append", opts=dict(title="Loss & Acc", legend=["loss", "acc"]))# 图像（注意通道顺序）img_bgr = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)img_rgb = img_bgr[:, :, ::-1]# BGR -> RGBimg = np.transpose(img_rgb, (2, 0, 1))# HWC -> CHWviz.image(img, win="sample", opts=dict(title="Sample Image"))# 文本日志viz.text("Hello, Visdom!", win="log", opts=dict(title="Log"))

  • 在浏览器访问：http://localhost:8097，选择环境demo查看。

四、模型结构与数据流可视化

• Netron：查看模型结构与张量形状
  • 安装：pip install netron
  • 使用：
    • 脚本方式：netron model.pth 或 netron.start("model.pth")（会打开浏览器）。
    • 命令行服务：netron model.pb --port 8080，访问 http://localhost:8080。
• torchinfo：结构化模型摘要
  • 安装：pip install torchinfo
  • 使用：

    from torchinfo import summaryimport torchvisionmodel = torchvision.models.resnet18()summary(model, input_size=(1, 3, 224, 224), device="cpu")

• PyTorchViz：计算图可视化
  • 安装：pip install torchviz
  • 使用：

    import torch, torchvisionfrom torchviz import make_dotmodel = torchvision.models.resnet18()x = torch.randn(1, 3, 224, 224)y = model(x)dot = make_dot(y, params=dict(model.named_parameters()))dot.render("resnet18_graph", format="pdf")# 生成PDF

• CNN内部可视化（进阶）：可进一步做卷积核、特征图与CAM/Grad-CAM可视化，用于理解网络学到了什么。

五、远程服务器与常见问题

• 远程访问
  • 服务器启动TensorBoard或Visdom后，在本地机器添加SSH隧道：
    • TensorBoard：ssh -L 6006:127.0.0.1:6006 user@server
    • Visdom：ssh -L 8097:127.0.0.1:8097 user@server
  • 然后在本地浏览器访问：http://localhost:6006 或 http://localhost:8097。
• 端口与防火墙
  • 确保服务器防火墙放行对应端口（如6006/8097），或改用未被占用的端口启动服务。
• 图像通道顺序
  • Visdom默认使用RGB，而OpenCV读取为BGR，显示前请做通道转换：img_rgb = img_bgr[:, :, ::-1]。
• 日志目录管理
  • 多实验建议分目录记录，如：runs/exp1/、runs/exp2/，便于TensorBoard对比。