本文介绍了 PyTorch 模型训练的九个优化技巧，涵盖数据加载、批量大小、梯度累积、计算图管理、GPU 训练、混合精度及多卡分布式训练等方面。通过 DataLoader 并行加载、设置 num_workers、增大 batch size 可减少 IO 瓶颈；梯度累积可在显存受限时模拟大批量；使用 .item() 避免计算图残留；混合精度训练可节省内存并加速计算；多 GPU 和分布式训练利用 DDP 实现高效并行。这些方法能有效提升训练速度与资源利用率。
在深度学习开发中，训练效率直接影响实验迭代速度和资源成本。许多开发者仍在使用 32 位精度计算或在单 GPU 上串行训练，这往往导致内存浪费和训练缓慢。随着硬件与框架的演进，通过混合精度、多卡并行及数据加载优化等手段，可以显著提升 PyTorch 模型的训练性能。
以下总结了提升 PyTorch 模型训练速度的 9 个核心技巧，主要基于 PyTorch-Lightning 库的最佳实践进行说明。PyTorch-Lightning 是建立在 PyTorch 之上的高层封装，提供了自动化训练功能，同时允许开发者完全控制关键模型组件。
使用 DataLoader 来加载数据是获得训练速度提升的最简单方法之一。传统的 h5py 或 numpy 文件存储方式已逐渐被更高效的流式加载取代。对于图像数据，直接使用 PyTorch 的 DataLoader；对于 NLP 数据，可参考 TorchText 库。
在 PyTorch-Lightning 中，无需显式编写训练循环，只需定义好 DataLoaders 和 Trainer，框架会自动调用它们。
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
dataset = MNIST(root='./data', train=True, download=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
for batch in loader:
x, y = batch
model.training_step(x, y)
在此示例中，首先创建数据集实例，然后使用 DataLoader 封装。shuffle=True 确保数据随机打乱，pin_memory=True 可加速 CPU 到 GPU 的数据传输。DataLoader 支持批量大小调整，可根据实际需求优化。
在 DataLoader 中，设置 num_workers 参数允许批量并行加载数据，从而减少 IO 瓶颈。默认情况下，数据在主进程加载，速度较慢。
# 慢：主进程加载
loader = DataLoader(dataset, batch_size=32, shuffle=True)
# 快：启用 4 个 worker 进程
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)
将 num_workers 设置为 CPU 核心数附近（如 4 或 8），可启用并行加载。但需注意，过多的 worker 可能导致资源竞争，应根据系统内存和 IO 能力调整。
增加批量大小（batch size）到硬件允许的最大范围是重要的优化策略。较大的 batch size 能带来以下好处：
挑战在于内存占用增加和学习率调整。通常增大 batch size 后需按比例增加学习率（线性缩放规则）。需确保硬件支持并相应调整超参数。
当显存不足以支撑大 batch size 时，梯度累积是一种模拟大批量的技术。通过多次前向传播累积梯度，再执行一次反向传播和优化步骤。
optimizer.zero_grad()
scaled_loss = 0
accumulated_steps = 4
for i in range(accumulated_steps):
out = model.forward()
loss = some_loss(out, y) / accumulated_steps
loss.backward()
scaled_loss += loss.item()
optimizer.step()
actual_loss = scaled_loss
在 PyTorch-Lightning 中，可直接设置 accumulate_grad_batches 参数：
trainer = Trainer(accumulate_grad_batches=4)
trainer.fit(model)
记录损失值时，应只存储数值而非整个计算图，以节省内存。使用 .item() 方法获取标量值。
# 错误：保留计算图副本
losses.append(loss)
# 正确：仅存储数值
losses.append(loss.item())
若直接 append tensor，会保留引用导致内存无法释放。.item() 提取 Python 浮点数，断开计算图连接。
GPU 训练需将模型和数据移至 GPU 设备。虽然 PyTorch-Lightning 可自动处理，但理解底层机制有助于调试。
model.cuda()
x = x.cuda()
out = model(x)
注意限制 CPU-GPU 数据传输次数，避免频繁复制。必要时调用 torch.cuda.empty_cache() 清理缓存，但需谨慎使用，因其会阻塞同步。
使用 16 位浮点数（FP16）可将内存占用减半，且现代 GPU 对 FP16 有专门加速单元。混合精度（Mixed Precision）指部分计算用 FP16，权重保持 FP32。
原生 PyTorch 可使用 AMP（Automatic Mixed Precision）：
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
with autocast():
output = model(input)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
在 PyTorch-Lightning 中，设置 precision=16 即可自动启用：
trainer = Trainer(precision=16)
trainer.fit(model)
多 GPU 训练主要有三种模式：
将模型复制到每个 GPU，分配不同批次数据。适合单机多卡。
model = DataParallel(model, device_ids=[0, 1, 2, 3])
out = model(x)
Lightning 中设置 gpus=[0, 1, 2, 3] 即可。
适用于超大模型无法放入单卡显存的情况。将编码器放在 GPU 0，解码器放在 GPU 1。
self.encoder.cuda(0)
self.decoder.cuda(1)
out = self.decoder(self.encoder(x))
结合上述两种方式，针对不同模块分配不同设备。
在分布式训练中，每个机器上的每个 GPU 都有模型副本，独立初始化并在数据分区上训练，随后同步梯度更新。使用 DistributedDataParallel (DDP) 可实现高效同步。
def main_process_entrypoint(gpu_nb):
dist.init_process_group("nccl", rank=gpu_nb, world_size=world)
torch.cuda.set_device(gpu_nb)
model = DistributedDataParallel(model, device_ids=[gpu_nb])
# 训练逻辑...
if __name__ == '__main__':
mp.spawn(main_process_entrypoint, nprocs=8)
在 Lightning 中，配置更为简单：
trainer = Trainer(gpus=8, accelerator='ddp')
trainer.fit(model)
通过合理配置数据加载、批量大小、精度类型及并行策略，可显著提升 PyTorch 训练效率。建议优先尝试增大 batch size 和使用混合精度，其次考虑多卡并行。在实际项目中，应根据硬件资源和模型规模灵活组合上述技巧。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online