ubuntu pytorch慢怎么办

在 Ubuntu 上遇到 PyTorch 运行慢 的情况，通常和 CPU / GPU 环境、版本匹配、系统配置 有关。下面按「最常见 → 进阶」逐步排查和优化。

一、先确认是不是在用 GPU（最常见原因）

1. 检查 PyTorch 是否识别到 GPU

import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.device_count())

✅ 正确情况：

True
1

❌ 如果 False：

• 你装的是 CPU 版 PyTorch
• 或 CUDA 驱动 / 版本不匹配

二、安装正确的 PyTorch（GPU 版）

Ubuntu + NVIDIA 显卡推荐安装方式

1️⃣ 先确认驱动

nvidia-smi

能看到显卡信息说明驱动 OK。

2️⃣ 根据 CUDA 版本安装 PyTorch

例如 CUDA 11.8：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

⚠️ 不要直接 pip install torch，那默认是 CPU 版！

三、CPU 版 PyTorch 很慢怎么办？

如果你只能用 CPU：

1. 使用并行 BLAS

pip install intel-openmp

或设置线程数：

import torch
torch.set_num_threads(8)  # 看你 CPU 核心数

2. 用 torch.compile（PyTorch 2.x）

model = torch.compile(model)

3. 批量处理（batch）

❌ 慢：

for x in data:
    model(x)

✅ 快：

model(batch_x)

四、Ubuntu 系统层面的优化

1. 关闭 CPU 节能模式

sudo apt install cpufrequtils
sudo cpufreq-set -g performance

2. 禁用 nouveau（NVIDIA 必须）

lsmod | grep nouveau

如有，禁用：

sudo nano /etc/modprobe.d/blacklist.conf
# 添加
blacklist nouveau

五、Docker / WSL 用户特别注意

WSL2

• 必须安装 NVIDIA CUDA for WSL
• 否则 PyTorch 只能 CPU

Docker

用官方镜像：

docker run --gpus all pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel

六、快速自检清单 ✅

• [ ] torch.cuda.is_available() == True
• [ ] 安装的是 cuXXX 版本
• [ ] batch size 合理
• [ ] 没有频繁 .cpu() / .numpy()
• [ ] 模型在 cuda() 上

七、如果你愿意，我可以更具体地帮你

你可以告诉我：

1. Ubuntu 版本
2. 是否有 NVIDIA 显卡
3. nvidia-smi 输出
4. torch.__version__
5. 是训练慢还是推理慢

我可以直接给你最优安装命令和配置方案。