ubuntu pytorch如何调试

在 Ubuntu 上调试 PyTorch 代码，常见方式可以分为日志/打印调试、Python 调试器、PyTorch 专用工具、性能与分布式调试几类。下面按实用程度给你一个清晰清单。

一、最基础：打印与日志（快速定位）

1. print / logging

import logging
logging.basicConfig(level=logging.INFO)
logging.info(f"x shape: {x.shape}")

2. 检查张量状态

print(x.shape, x.dtype, x.device)
print(torch.isnan(x).any())
print(x.min(), x.max())

二、Python 调试器（强烈推荐）

1. pdb（命令行）

import pdb
pdb.set_trace()

运行后可用：

• n 下一步
• s 进入函数
• c 继续
• l 看代码
• p x 打印变量

2. ipdb（更好用）

pip install ipdb

import ipdb
ipdb.set_trace()

三、IDE 调试（最舒服）

PyCharm（Ubuntu 常见）

• Run → Edit Configurations
• 选择 Python 解释器（最好是 conda/venv）
• 打断点 → Debug

VS Code（轻量推荐）

• 安装 Python 插件
• 创建 .vscode/launch.json

{
  "version": "0.2.0",
  "configurations": [
    {
      "name": "Python: PyTorch",
      "type": "python",
      "request": "launch",
      "program": "train.py",
      "console": "integratedTerminal"
    }
  ]
}

四、PyTorch 专属调试技巧

1. 自动检测 NaN / Inf

torch.autograd.set_detect_anomaly(True)

配合：

with torch.autograd.detect_anomaly():
    loss.backward()

2. 钩子（hook）调试

def hook(module, input, output):
    print(module, output.shape)

model.layer1.register_forward_hook(hook)

3. torchinfo 看模型结构

pip install torchinfo

from torchinfo import summary
summary(model, input_size=(1,3,224,224))

五、性能调试（慢在哪里）

1. 时间分析

import time
start = time.time()
output = model(x)
print(time.time() - start)

2. torch.utils.bottleneck

python -m torch.utils.bottleneck train.py

3. py-spy（不修改代码）

pip install py-spy
sudo py-spy top --pid <pid>

六、分布式 / GPU 调试

1. 确认 GPU 状态

nvidia-smi

2. 单卡先调通

CUDA_VISIBLE_DEVICES=0 python train.py

3. Distributed 调试

import os
print("RANK:", os.environ["RANK"])

七、常见坑提醒（Ubuntu 特有）

• CUDA 版本不匹配

nvcc --version
python -c "import torch; print(torch.version.cuda)"

• conda / 系统 Python 混用

which python
which torch

• 权限问题

sudo chmod -R 777 /tmp

八、推荐组合（新手友好）

✅ VS Code + ipdb + torch.autograd.set_detect_anomaly

如果你愿意，可以告诉我：

• 你是 训练 / 推理 / 分布式？
• 用的是 conda 还是系统 Python？
• 报错类型（NaN / 卡死 / OOM）？

我可以直接给你针对性的调试步骤。