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 SummaryWriter
    import torch, torchvision
    
    writer = 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, time
    
    viz = 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与acc
    for step in range(100):
        loss = abs(np.random.randn()) + 1
        acc  = 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 -> RGB
    img = np.transpose(img_rgb, (2, 0, 1))  # HWC -> CHW
    viz.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 summary
    import torchvision
    
    model = torchvision.models.resnet18()
    summary(model, input_size=(1, 3, 224, 224), device="cpu")
    

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

    import torch, torchvision
    from torchviz import make_dot
    
    model = 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对比。