Ubuntu环境下如何调试PyTorch代码

Ubuntu下调试PyTorch代码的实用流程

一 环境准备与快速检查

• 建议使用Anaconda或venv创建隔离环境，避免依赖冲突。
• 安装匹配你CUDA版本的PyTorch（CPU版亦可）：
  pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
  或在conda中：conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
• GPU环境快速自检：
  • 查看驱动与CUDA：nvidia-smi
  • 进入Python：
    import torch; print(torch.version, torch.cuda.is_available())
• 若遇到环境或链接问题，设置环境变量：在~/.bashrc或~/.zshrc中加入
  export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-“$(dirname $(which conda))/…/”}
  并执行 source ~/.bashrc 或 source ~/.zshrc 使配置生效。

二 应用层调试方法

• Python调试器
  • pdb：在需要处插入 import pdb; pdb.set_trace()，支持断点、单步、打印变量等。
  • ipdb：增强版pdb，提供彩色提示与补全，使用 import ipdb; ipdb.set_trace()。
• IDE图形化调试
  • PyCharm：在行号左侧点下断点，点击“虫子”图标启动调试，可在调试控制台查看变量与调用栈。
  • VSCode：安装Python扩展，生成或编辑 .vscode/launch.json，配置 program、args、env 等后点击“开始调试”。
• 交互式与日志
  • Jupyter Notebook：在异常后使用 %debug 进入事后调试；也可在单元中插入断点。
  • 日志与断言：使用 logging 记录关键变量与流程；用 assert 快速校验形状、范围与设备一致性。
• 专用工具
  • torchsnooper：函数装饰器 @torchsnooper.snoop() 自动打印张量的形状、dtype、device、requires_grad 等，定位张量相关错误非常高效。

三 训练与性能问题定位

• 异常与梯度检查
  • 在训练开始前启用：torch.autograd.set_detect_anomaly(True)，可在反向传播中更准确地定位产生 NaN/Inf 的操作。
• 性能瓶颈分析
  • 使用 PyTorch Profiler 定位算子耗时与GPU利用率：
    with torch.profiler.profile(on_trace_ready=torch.profiler.tensorboard_trace_handler(“trace.pt”)) as prof:
    for step, data in enumerate(trainloader, 0):
    inputs, labels = data[0].to(device), data[1].to(device)
    outputs = model(inputs); loss = criterion(outputs, labels)
    optimizer.zero_grad(); loss.backward(); optimizer.step()
    • 在 TensorBoard 中查看 trace：tensorboard --logdir trace.pt
• 训练过程可视化
  • 使用 torch.utils.tensorboard 记录损失、学习率、梯度与权重直方图，辅助排查发散与收敛问题。

四 GPU与多进程常见故障排查

• 环境与版本
  • 确认 CUDA/cuDNN 与 PyTorch 版本匹配；用 nvidia-smi 检查驱动与显存；在Python中确认 torch.cuda.is_available()。
• 数据与模型设备一致性
  • 排查“Tensor on CPU vs. GPU”不匹配：统一使用 .to(device)，打印关键张量的 device 与 dtype。
• 分布式训练
  • 优先用单进程/单卡复现问题，再逐步扩展到多卡；检查初始化、端口占用与 NCCL 超时设置。
• 确定性复现
  • 固定随机种子：
    import torch, numpy, random
    torch.manual_seed(0); torch.cuda.manual_seed_all(0)
    numpy.random.seed(0); random.seed(0)
    torch.backends.cudnn.deterministic=True; torch.backends.cudnn.benchmark=False

五 深入源码级调试

• 适用场景：需要定位 PyTorch C++/CUDA 算子或底层库问题时。
• 基本步骤
  • 使用 Anaconda 创建环境并拉取源码：
    conda create -n torch-dev python=3.x
    conda activate torch-dev
    git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch && git submodule sync && git submodule update --init --recursive
  • 带调试信息编译（示例关闭部分特性以加速首次构建）：
    export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-“$(dirname $(which conda))/…/”}
    DEBUG=1 USE_DISTRIBUTED=0 USE_MKLDNN=0 USE_CUDA=0 BUILD_TEST=0 USE_FBGEMM=0 USE_NNPACK=0 USE_QNNPACK=0 python setup.py develop
  • 远程调试：在服务器上以调试模式运行脚本，使用 VSCode Remote-SSH 连接；结合 gdb/VS Code C++ 调试器附加到相关进程，或在Python层配合 pdb/ipdb 联动定位。