Ubuntu下PyTorch的调试工具有哪些

Ubuntu下PyTorch调试工具全景

一 基础调试与断点工具

• pdb / ipdb：在代码中插入断点，逐行执行、查看变量与栈帧。示例：import pdb; pdb.set_trace()；或在异常后使用 %debug（Notebook）进入事后调试。适合快速定位逻辑错误与形状不匹配等问题。
• IDE 图形化调试：PyCharm、VSCode 支持设置断点、条件断点、观察表达式、单步执行等，适合工程化开发与复杂调用栈排查。
• 交互式环境：IPython.embed() 在任意位置“就地”打开交互式会话；Jupyter Notebook 的 %debug 便于单元级问题定位。
• 日志与断言：使用 logging 输出关键变量与执行路径；用 assert 快速校验张量形状、取值范围与设备一致性。

二 PyTorch专属与异常检测

• torch.autograd.set_detect_anomaly(True)：在反向传播中检测数值异常（如 NaN/Inf、梯度异常），对定位训练发散、loss 为 NaN 等问题非常有效。
• TorchSnooper：自动追踪函数内张量的形状、dtype、device、requires_grad等关键信息，装饰器用法简单，适合快速“看清”数据在每层的变化。
• VizTracer：低开销的代码跟踪/性能分析工具，可生成时间线；对 PyTorch 可用 --log_torch 或 log_torch=True 记录原生调用与 GPU 事件，便于定位性能瓶颈与执行路径。

三 性能与计算图可视化

• TensorBoard + torch.utils.tensorboard：记录并可视化标量（如 loss/accuracy）、直方图、图像、图结构等，便于训练过程“看得见”的调试与对比。
• torch.autograd.profiler：结合 profile 与 record_function 做算子级性能剖析，定位训练/推理热点与设备利用率问题。
• 模型结构与计算图：
  • Torchinfo：一键输出模型各层参数量、输出形状、内存占用等摘要，快速发现维度不匹配与参数规模异常。
  • PyTorchviz：基于 Graphviz 绘制前向计算图，直观查看节点依赖与张量流向。
  • Netron：可视化模型结构（支持多种格式），便于与部署端对齐模型定义与层次结构。

四 GPU与系统层面的排查

• nvidia-smi：查看 GPU 利用率、显存占用、驱动/CUDA 版本与健康状态，辅助判断是否存在 OOM、驱动不匹配或资源争用。
• cProfile：Python 内置CPU 性能分析，用于定位数据加载、预处理、训练循环等非 GPU 瓶颈。
• 环境与版本校验：确认 PyTorch、CUDA/cuDNN 版本匹配；必要时在 shell 配置 CMAKE_PREFIX_PATH 等环境变量以避免构建/加载问题。

五 实用组合与建议

• 逻辑/数值错误：pdb/VSCode 断点 + assert + torch.autograd.set_detect_anomaly(True) + TorchSnooper，快速定位张量形状、设备与梯度异常。
• 性能瓶颈：torch.autograd.profiler 定位算子热点，VizTracer --log_torch 查看调用时间线，TensorBoard 对比多轮指标与资源曲线。
• 模型结构/数据流：Torchinfo 做层级摘要，PyTorchviz 看计算图依赖，Netron 校验导出模型结构一致性。
• GPU/环境：nvidia-smi + cProfile + 版本/环境变量检查，先排除资源与驱动问题，再回到模型与训练逻辑。