ubuntu pytorch怎么加速训

在 Ubuntu 上加速 PyTorch 训练，通常从硬件利用、软件配置、代码优化三个层面入手。下面按实用程度给你一份清单。

一、确认并开启 GPU 加速（最重要）

1. 检查 CUDA 是否可用

import torch
print(torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.cuda.get_device_name(0))

如果返回 False：

• 确认装的是 GPU 版 PyTorch

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

（cu118 对应 CUDA 11.8，按你驱动版本选）

2. 把模型和数据放到 GPU

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
inputs, labels = inputs.to(device), labels.to(device)

二、使用 cuDNN 自动优化

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

✅ 适合输入尺寸固定的情况
❌ 输入尺寸变化频繁时不要开

三、使用混合精度（AMP，非常有效）

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for x, y in dataloader:
    x, y = x.cuda(), y.cuda()
    with autocast():
        out = model(x)
        loss = criterion(out, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

✅ 显存 ↓ 30%~50%
✅ 速度 ↑ 1.5~2x（V100 / A100 / 30 系以上）

四、DataLoader 加速（常被忽略）

torch.utils.data.DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,        # CPU 核心数左右
    pin_memory=True,      # GPU 训练必开
    prefetch_factor=4
)

✅ pin_memory=True 能明显提速
✅ num_workers 不是越大越好（一般 = CPU 核数）

五、多 GPU / 分布式训练

单机多卡（简单）

model = torch.nn.DataParallel(model)

推荐（更快更稳）

model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[local_rank]
)

配合：

torchrun --nproc_per_node=2 train.py

六、Ubuntu 系统层面优化

1. 更新 NVIDIA 驱动

nvidia-smi

建议驱动 ≥ 525（新卡用 535+）

2. 关闭 CPU 节能

sudo cpupower frequency-set -g performance

3. 使用 NVMe / 高速磁盘

数据读取是常见瓶颈

七、代码级提速技巧

• 用 torch.compile（PyTorch 2.x）

model = torch.compile(model)

• 减少 .item()、.cpu() 频繁调用
• 用 in-place 操作（谨慎）
• 避免 Python 循环，多用向量化

八、推荐顺序（立竿见影）

1. ✅ 确认 GPU 可用
2. ✅ 开 cudnn.benchmark
3. ✅ 用 AMP
4. ✅ DataLoader 调好
5. ✅ 多卡 DDP
6. ✅ torch.compile

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

• GPU 型号（如 3090 / A100）
• PyTorch 版本
• 是不是训练 CNN / Transformer
• 单卡还是多卡

我可以直接帮你给一份可复制的训练加速模板。