PyTorch CUDA自动混合精度训练详解

"torch.cuda.amp - 自动混合精度详解" 自动混合精度（Automatic Mixed Precision, AMP）是PyTorch库中的一个功能，用于在训练深度学习模型时提高计算效率和内存利用率，尤其是在GPU上。它通过使用半精度（FP16）数据类型来加速计算，同时利用全精度（FP32）来保持数值稳定性。torch.cuda.amp模块提供了实现这一功能的工具。 在给定的代码片段中，可以看到如何在训练过程中启用和管理AMP： 1. **定义scaler对象**：`scaler=torch.cuda.amp.GradScaler()` 是AMP的核心组件，它负责动态调整比例因子以保持梯度的数值稳定。只有在PyTorch版本大于等于1.6时，才能使用此功能。 2. **前向传播**：使用`autocast`上下文管理器将前向传播操作置于半精度环境中。`from torch.cuda.amp import autocast as autocast`导入autocast函数，然后`with autocast():`包围模型的前向传播代码，使得模型内部的计算在FP16模式下进行。 3. **反向传播**：在AMP模式下，梯度的计算需要特别处理。如果PyTorch版本大于等于1.6，使用`scaler.scale(loss).backward()`。`scaler.scale(loss)`会将损失乘以当前的比例因子，然后调用`backward()`进行反向传播。这样可以确保在半精度下仍能保持梯度的数值稳定性。 4. **梯度剪切**：无论PyTorch版本如何，都需要对梯度进行剪切，以防止梯度过大导致数值溢出。这里使用`torch.nn.utils.clip_grad_norm_`对梯度进行裁剪，限制最大范数。在小于1.6的版本中，`amp.master_params(optimizer)`返回需要裁剪的主参数列表，而在1.6及以上版本，直接对模型的`parameters()`进行裁剪。 5. **优化器更新**：在AMP环境下，优化器的步进操作也需要特别处理。在1.6及以上版本，使用`scaler.step(optimizer)`和`scaler.update()`分别执行优化器的步进和更新比例因子的操作。而在较旧版本中，直接调用`optimizer.step()`。 6. **模型保存**：在AMP模式下，如果使用了1.6以下的PyTorch版本，需要保存`amp.state_dict()`以恢复AMP状态。在1.6及以上版本，AMP的状态不再需要单独保存，只需保存模型和优化器的状态。 通过使用自动混合精度，可以显著加快训练速度，尤其对于大规模模型和计算密集型任务，同时保持模型的准确度。不过需要注意的是，启用AMP后可能需要调整学习率、梯度裁剪阈值等超参数，以适应半精度环境。