Python AI 实战（12/20）：AI 模型训练太慢怎么加速？batch size、混合精度、多卡方案

Python AI 实战问答速查手册 · 第 12/20 问

所属板块：模型训练与调优

先找瓶颈再看方案：GPU 利用率低可能来自数据加载、CPU 运算、同步等待或 batch 过小；用分段计时或 profiler 判断，不能只凭一个利用率数值下结论。

方案一：单卡加速（先做这些，性价比最高）

1. 混合精度训练（AMP）：让适合的算子使用半精度，并对 fp16 训练进行梯度缩放。收益取决于硬件和模型。以下是 CUDA 训练步骤片段，假设 model、optimizer、loss_fn、train_loader 已创建，模型及数据在 CUDA 上：

import torch

scaler = torch.amp.GradScaler("cuda")  # 在训练循环外创建
for x, y in train_loader:
  x, y = x.to("cuda"), y.to("cuda")
  optimizer.zero_grad(set_to_none=True)
  with torch.autocast(device_type="cuda", dtype=torch.float16):
    prediction = model(x)
    loss = loss_fn(prediction, y)  # 标量损失
  scaler.scale(loss).backward()
  scaler.step(optimizer)
  scaler.update()

1. 调大 batch size：GPU 并行能力吃满；调整后重新验证学习率和泛化表现，线性缩放并非适用于所有任务。
2. 数据加载优化：DataLoader(..., num_workers=4, pin_memory=True, persistent_workers=True)，把数据预处理放在子进程；瓶颈在 CPU 预处理时考虑在 GPU 上做或预缓存。
3. 减少 CPU-GPU 同步：训练循环里少写 loss.item()、.cpu()、print（每次同步都会让 GPU 等 CPU）。
4. torch.compile：PyTorch 2.x 一行 model = torch.compile(model)，需要把编译预热开销与稳定运行速度分别测量。
5. 梯度累积：显存不够又想要大 batch 时，多步小 batch 累积梯度再更新。

方案二：多卡训练

• DDP（DistributedDataParallel）：官方推荐，多进程每卡一份数据，同步梯度；加速效果取决于计算与通信比例，train.py 还需初始化进程组、包装 DDP 并正确划分数据：

torchrun --nproc_per_node=4 train.py

• 大模型（单卡放不下）：用 DeepSpeed / FSDP（ZeRO 分片，把模型参数分摊到多卡）。

• 注意：小模型多卡收益低（通信开销占比大），先把单卡榨干再上多卡。

方案三：模型层面

• 换更小/更高效的模型结构（如 MobileNet 替代 ResNet-152、蒸馏得到的小模型）。

• 模型量化（训练后 INT8 量化）主要提速推理，训练阶段收益小。

方案四：云上省钱思路：先本地/小卡验证代码，再上大卡跑全量；按 GPU 小时计费，跑之前先测一下单 step 时间估算总时长。