Ubuntu上如何调试PyTorch程序

在Ubuntu上调试PyTorch程序可按以下步骤进行：

一、环境配置

1. 安装基础工具
   • 使用Miniconda管理虚拟环境：

     wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
     bash Miniconda3-latest-Linux-x86_64.sh
     conda create -n pytorch_env python=3.8
     conda activate pytorch_env
     

   • 安装PyTorch（含CUDA支持，需匹配GPU架构）：

     conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch
     

2. 验证安装

   import torch
   print(torch.__version__, torch.cuda.is_available())  # 检查版本和CUDA是否可用
   

二、调试工具与方法

1. 交互式调试
   • IPython/Jupyter Notebook：直接运行代码，支持逐步执行和变量查看。
   • IDE调试（如PyCharm/VSCode）：
     • 设置断点：点击代码行左侧空白处。
     • 启动调试：点击“Debug”按钮，通过控制台查看变量、单步执行（F5/F10）。
2. 代码级调试
   • pdb断点：在代码中插入import pdb; pdb.set_trace()，程序暂停后可使用n（下一步）、s（进入函数）、c（继续）等命令。
   • 日志记录：
     • 使用Python logging模块记录关键信息。
     • 用print()输出变量值（简单场景适用）。
3. 性能与可视化工具
   • TensorBoard：记录损失、梯度等指标，可视化训练过程：

     from torch.utils.tensorboard import SummaryWriter
     writer = SummaryWriter('runs/experiment')
     writer.add_scalar('Loss/train', loss.item(), epoch)
     

   • PyTorch专用工具：
     • torch.autograd.set_detect_anomaly(True)：检测梯度异常。
     • torch.autograd.profiler：分析计算性能。

三、进阶调试技巧

• 单元测试：用unittest或pytest编写测试用例，验证模块功能。
• 混合精度训练：用torch.cuda.amp减少内存占用，加速计算。
• 分布式调试：多GPU场景下，使用torch.distributed模块时需检查进程同步和通信。

四、常见问题排查

• CUDA错误：检查驱动版本与CUDA Toolkit是否兼容，通过nvidia-smi查看驱动状态。
• 内存不足：减小batch size，或使用torch.utils.checkpoint进行梯度检查点优化。

参考资料：