混合精度计算（Mixed Precision Training）是一种在深度学习训练过程中同时使用单精度浮点数（FP32）和半精度浮点数（FP16）进行计算的方法。其主要目的是在不显著影响模型精度的前提下，提高计算效率和减少显存占用。
混合精度计算的优势
- 提高计算效率：FP16计算在现代GPU（如NVIDIA的Volta、Turing和Ampere架构）上可以显著加速，因为这些GPU对FP16运算进行了优化。
- 减少显存占用：FP16数据占用的显存空间是FP32的一半，这意味着可以在相同的显存中处理更大的模型或更大的批量数据。
实现混合精度计算的关键组件
- 自动混合精度（AMP）：PyTorch提供了
torch.cuda.amp模块来简化混合精度训练的实现。AMP会自动管理FP16和FP32的转换，确保计算的稳定性和精度。 - 梯度缩放（Gradient Scaling）：由于FP16的动态范围较小，梯度值可能会变得非常小，从而导致数值下溢。梯度缩放通过在反向传播之前将梯度值放大，在更新权重之前再缩小，来避免这个问题。
在这个示例中，我们使用了torch.cuda.amp.autocast来自动进行混合精度计算，并使用torch.cuda.amp.GradScaler来缩放损失值，从而避免数值不稳定性。这样可以显著减少显存占用，同时加速训练过程。
代码如下所示
import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler
# 定义一个简单的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 初始化模型、损失函数和优化器
model = SimpleModel().cuda()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 创建GradScaler对象
scaler = GradScaler()
# 模拟训练数据
inputs = torch.randn(32, 10).cuda()
targets = torch.randn(32, 1).cuda()
# 训练循环
for epoch in range(10):
model.train()
optimizer.zero_grad()
# 使用autocast进行前向和后向传播
with autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
# 使用GradScaler进行反向传播和优化
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Epoch [{epoch+1}/10], Loss: {loss.item():.4f}")核心代码：
from torch.cuda.amp import autocast, GradScaler
……
# 使用autocast进行前向和后向传播
with autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
# 使用GradScaler进行反向传播和优化
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
