linux pytorch怎么加速

在 Linux + PyTorch 环境下，加速训练/推理通常可以从硬件、系统、PyTorch 配置、代码层面四个方向入手。下面给你一份实用清单（按性价比排序）。

一、确认并充分利用 GPU（最关键）

1. 确认 GPU 可用

nvidia-smi
python -c "import torch; print(torch.cuda.is_available())"

2. 使用 GPU

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

3. 使用 torch.cuda.amp（混合精度）

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

✅ 通常可提速 30%~2x，显存减半

二、DataLoader 加速（常被忽略）

1. 正确使用 num_workers

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

• num_workers ≈ CPU 核心数
• pin_memory=True（GPU 训练必开）

2. 数据预处理前置

• 用 NVMe SSD
• 预处理缓存（.npy / .pt）
• 或用 WebDataset / LMDB

三、PyTorch 自身加速选项

1. torch.backends 开启

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

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

2. 使用 torch.compile（PyTorch 2.0+）

model = torch.compile(model)

✅ A100 / 3090 上常提速 10%~30%

四、模型 & 计算优化

1. 使用高效模型结构

• ConvNeXt / EfficientNet / ViT
• 减少 FC 层
• 使用 depthwise conv

2. 减少 Python 开销

• 避免 for 循环
• 使用 torch.nn.functional
• 减少 .item() / .cpu() 调用

五、分布式 & 多卡

1. 单机多卡（DDP）

torchrun --nproc_per_node=4 train.py

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

✅ 接近线性加速

2. 多机（Slurm / MPI）

适合大模型训练

六、系统级优化（Linux）

1. 禁用 CPU 节能

sudo cpupower frequency-set -g performance

2. 提高文件句柄

ulimit -n 65535

3. 使用合适的驱动 & CUDA

• CUDA ≥ 11.8
• 驱动匹配 PyTorch 版本

七、推理加速（如果你关心预测）

• torch.jit.trace
• TensorRT
• ONNX + ONNXRuntime
• torch.cuda.graphs

快速 checklist ✅

• [ ] GPU 可用
• [ ] model.to(cuda)
• [ ] DataLoader(num_workers, pin_memory)
• [ ] cudnn.benchmark=True
• [ ] amp / torch.compile
• [ ] DDP（多卡）

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

• ✅ 是 训练还是推理
• ✅ GPU 型号（如 3090 / A100）
• ✅ 模型类型（CNN / Transformer / LLM）

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