ubuntu pytorch怎样加速

在 Ubuntu 上加速 PyTorch，通常可以从硬件加速、软件配置、代码优化三个层面来做。下面按实用程度给你一套清单。

一、使用 GPU（CUDA）加速（最重要）

1. 确认有 NVIDIA 显卡

nvidia-smi

能看到 GPU 信息说明驱动已装好。

2. 安装合适版本的 PyTorch（CUDA 版）

不要用 pip 默认 CPU 版。

# 例如 CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

验证：

import torch
print(torch.cuda.is_available())  # True

3. 把模型和数据放到 GPU

device = torch.device("cuda")
model.to(device)
x = x.to(device)

二、Ubuntu 系统级优化

1. 更新 NVIDIA 驱动

ubuntu-drivers devices
sudo ubuntu-drivers autoinstall

2. 开启性能模式

sudo apt install cpufrequtils
sudo cpufreq-set -g performance

3. 禁用 Nouveau（如仍存在）

echo "blacklist nouveau" | sudo tee /etc/modprobe.d/blacklist-nouveau.conf
sudo update-initramfs -u

三、PyTorch 运行参数加速

1. 使用 cuDNN 自动优化

torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True

2. 使用混合精度（AMP）

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    output = model(x)
    loss = loss_fn(output, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

四、DataLoader 加速

torch.utils.data.DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,        # CPU 核心数
    pin_memory=True,
    prefetch_factor=4
)

五、多 GPU / 分布式（进阶）

1. 单机多卡

model = torch.nn.DataParallel(model)

2. 推荐方式（DDP）

torchrun --nproc_per_node=2 train.py

六、其他加速技巧

• 使用 TorchScript / torch.compile（PyTorch 2.x）

model = torch.compile(model)

• 减少 Python 循环
• 使用 inplace=True（谨慎）
• 使用更高效算子（如 FlashAttention）

七、CPU 用户怎么办？

如果是无 GPU：

• 安装 Intel MKL 版 PyTorch
• 使用 torch.set_num_threads(N)
• 使用 ONNX / TensorRT（推理）

如果你能告诉我：

• Ubuntu 版本
• 是否有 GPU（型号）
• 训练还是推理
• PyTorch 版本

我可以给你针对性的加速方案。