趋近智
训练大型扩散模型,以其复杂的架构(如改进的 U-Net 或 Transformer)而闻名,需要大量的计算资源和时间。随着模型复杂度的提高和数据集的扩大,优化训练过程变得必不可少,这不仅是为了速度,也是为了在硬件限制下实现可行性。混合精度训练是一种有效技术,通过策略性地使用低精度浮点数进行计算来应对这些挑战,从而带来显著的加速和内存消耗的减少。
深度学习 (deep learning)模型通常使用 32 位单精度浮点数 (FP32) 来存储权重 (weight)、激活和梯度。每个 FP32 数字使用 32 位内存。混合精度训练引入了低精度格式:
FP32、FP16 和 BF16 浮点格式的位分配、范围和精度特征的比较。
简单地将所有操作切换到 FP16 可能会导致数值不稳定。有效的混合精度训练将低精度计算与保持精度和稳定性的方法相结合,通常由深度学习 (deep learning)框架自动完成:
主流深度学习 (deep learning)框架提供了便捷的抽象来启用混合精度:
torch.cuda.amp(自动混合精度)。主要组件是 torch.autocast,用于在其上下文 (context)中自动转换操作,以及 torch.cuda.amp.GradScaler,用于管理损失缩放。tf.keras.mixed_precision。你设置一个全局策略(例如 mixed_float16 或 mixed_bfloat16),并用 LossScaleOptimizer 包装优化器。以下是一个使用 PyTorch amp 的简化示例:
import torch
# scaler 通常在训练循环外部初始化一次
scaler = torch.cuda.amp.GradScaler(enabled=True) # 启用 AMP
# 在训练循环内部:
optimizer.zero_grad()
# 对前向传播(模型执行和损失计算)使用 autocast
# 自动为符合条件的操作选择 FP16/BF16
with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=True):
# 假设模型接收带噪声图像和时间步
predicted_noise = model(noisy_images, timesteps)
loss = loss_fn(predicted_noise, target_noise)
# 缩放损失。在缩放后的损失上调用 backward() 以生成缩放后的梯度。
scaler.scale(loss).backward()
# scaler.step() 首先对优化器分配的参数的梯度进行反缩放。
# 如果梯度不是 inf/NaN,则调用 optimizer.step()。
# 否则,optimizer.step() 被跳过。
scaler.step(optimizer)
# 更新下一次迭代的缩放因子。
scaler.update()
# 继续训练循环...
扩散模型的具体考量:
GradScaler 的缩放因子很有必要。如果缩放因子频繁下降,可能表明存在数值问题。确保涉及时间嵌入 (embedding)或归一化 (normalization)层的操作能正确处理精度变化;框架通常会处理此问题,但自定义层可能需要检查。优点:
权衡:
在实践中,混合精度训练是一种标准且非常有效的技术,用于加速大型扩散模型的开发和部署。它使得研究人员和工程师能够更快地迭代和训练更强大的模型,通过高效地使用现代硬件加速器。在训练大型扩散模型时,采用混合精度通常不只是一种优化,而是一种必要。
这部分内容有帮助吗?
torch.cuda.amp和GradScaler实现混合精度训练的官方指南。tf.keras.mixed_precision实现混合精度训练的官方指南。© 2026 ApX Machine Learning用心打造