linux pytorch如何提速训练

在 Linux + PyTorch 环境下，要从训练速度上“榨干”性能，通常可以从 硬件、系统、PyTorch 配置、模型与数据 四个层面入手。下面按优先级和可操作性给你一个实用清单。

一、硬件与系统层面（基础）

1. 使用 GPU 并正确配置

nvidia-smi

确认：

• CUDA 可用
• GPU 利用率高（训练时应 >80%）

PyTorch 使用 GPU：

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

2. 安装正确的 PyTorch 版本

用 CUDA 对应版本 的 PyTorch（不要只用 pip 默认 CPU 版）：

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

3. 开启 cuDNN 自动优化

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

⚠️ 输入尺寸固定时效果最好

二、PyTorch 训练提速核心手段（最常用）

4. 使用 DataLoader 多进程加载

DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,      # 通常 = CPU 核心数
    pin_memory=True     # GPU 训练必开
)

5. 使用 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()

✅ 显存↓，速度↑（通常 1.5–2×）

三、并行与分布式（进阶）

6. 单机多卡：DataParallel（简单但不推荐）

model = torch.nn.DataParallel(model)

7. 单机多卡：DistributedDataParallel（推荐）

torchrun --nproc_per_node=2 train.py

torch.distributed.init_process_group("nccl")
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

✅ 比 DP 快、更稳定

四、模型与训练策略优化

8. 增大 batch size（配合 AMP）

• 更大 batch → 更高 GPU 利用率
• 必要时调整学习率

9. 使用高效优化器

• torch.optim.AdamW
• 或 bitsandbytes（8-bit optimizer）
• 或 torch.optim.LAMB（大 batch）

10. 减少不必要操作

• 去掉 .item() / .cpu() 频繁调用
• 减少 print
• 不保存每步 loss 到 list

五、Linux 系统级优化（高阶）

11. 锁定 CPU 频率（避免降频）

sudo cpufreq-set -g performance

12. 关闭 SWAP（防止卡顿）

sudo swapoff -a

13. 使用 NVMe SSD 存放数据

数据读取速度直接影响训练

六、快速自检清单 ✅

七、典型加速效果参考

• AMP：+30%~100%
• DDP：接近线性加速
• 数据加载优化：+20%~50%

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

• 单卡 / 多卡？
• 模型类型（CNN / Transformer / LLM）？
• 数据规模？

我可以直接给你 可复制的训练加速模板代码。