ubuntu pytorch怎样加速训练

在 Ubuntu 上加速 PyTorch 训练，可以从硬件、软件环境、代码层面三个方向入手。下面给你一套实用清单（从最重要到进阶）。

一、硬件 & 驱动（最基础）

1. 使用 GPU（CUDA）

nvidia-smi

确认：

• 有 NVIDIA 显卡
• 驱动正常

安装匹配驱动：

ubuntu-drivers devices
sudo ubuntu-drivers autoinstall

2. 安装正确版本的 CUDA + cuDNN

PyTorch 官方推荐方式（最稳）：

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

验证：

import torch
print(torch.cuda.is_available())
print(torch.version.cuda)

二、PyTorch 训练加速（核心）

3. 使用 .to('cuda')

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)
inputs = inputs.to(device)

4. 使用 DataParallel 或 DistributedDataParallel

单机多卡（简单）

model = torch.nn.DataParallel(model)

多卡高效（推荐）

torchrun --nproc_per_node=2 train.py

model = torch.nn.parallel.DistributedDataParallel(model)

5. 开启 cuDNN 自动优化

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

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

model = torch.compile(model)

在 Ubuntu + GPU 上通常能加速 10–30%

7. 使用混合精度（AMP）

from torch.cuda.amp import autocast, GradScaler

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

三、数据加载加速

8. DataLoader 优化

DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,
    pin_memory=True,
    prefetch_factor=4
)

9. 数据预处理加速

• 用 opencv / PIL + cache
• SSD / NVMe 硬盘
• 避免 Python 循环

四、系统 & 环境优化

10. 使用 conda / venv（避免系统冲突）

conda create -n torch python=3.10
conda activate torch

11. 升级内核 & 关闭 CPU 节能

sudo apt update
sudo apt install linux-generic

五、进阶（可选）

• TensorRT 推理加速
• DeepSpeed（大模型）
• NVIDIA Apex
• 使用 torch.utils.benchmark

快速检查清单 ✅

• [ ] torch.cuda.is_available() == True
• [ ] cudnn.benchmark = True
• [ ] 使用 torch.compile
• [ ] 使用 AMP
• [ ] num_workers > 0
• [ ] 多卡用 DDP

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

• Ubuntu 版本
• GPU 型号
• 是否多卡
• 训练的是 CNN / Transformer / LLM

我可以给你针对性的配置示例。