PyTorch自动混合精度训练

torch.cuda.amp.GradScaler 是 PyTorch 中的一个用于自动混合精度(Automatic Mixed Precision, AMP)训练的工具。AMP 允许在训练深度学习模型时动态切换浮点数的精度(例如,使用半精度浮点数 float16 而非 float32),以减少显存占用和加速计算,同时保持模型的精度。

1. GradScaler 的作用

在混合精度训练中,模型的某些部分以半精度(float16)计算,而其他部分仍然以全精度(float32)计算。使用 float16 进行计算可以显著提高计算速度和减少显存占用,但也可能导致数值不稳定或梯度下溢(gradient underflow)。GradScaler 通过动态缩放损失值来缓解这些问题,并在反向传播过程中对缩放后的梯度进行适当调整,确保训练过程稳定。

2. 混合精度训练的基本步骤

2. 1. 初始化 GradScaler
scaler = torch.cuda.amp.GradScaler()
2. 2. 在前向传播中使用 autocast 上下文管理器

在模型的前向传播中,使用 torch.cuda.amp.autocast 上下文管理器将部分计算切换到半精度。

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

在 PyTorch 中,autocast是一个用于自动混合精度训练的上下文管理器。

2. 3 使用 scaler.scale 缩放损失并反向传播

在计算损失并调用 backward() 前,通过 scaler.scale() 对损失进行缩放。

scaler.scale(loss).backward()
2. 4 使用 scaler.step 进行优化器更新

使用 scaler.step() 来执行优化器的 step() 操作。

scaler.step(optimizer)

2.5 调用 scaler.update : 通过 scaler.update() 来更新缩放因子,并根据需要调整精度。

scaler.update()

3. 完整的示例代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast

# 假设你有一个简单的模型
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.SGD(model.parameters(), lr=0.001)

# 初始化 GradScaler
scaler = GradScaler()

# 假设你有输入和目标
inputs = torch.randn(64, 10).cuda()
targets = torch.randn(64, 1).cuda()

# 训练循环中的一次前向和反向传播
for epoch in range(10):
    optimizer.zero_grad()

    # 前向传播,使用 autocast 进行混合精度计算
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # 反向传播,使用 scaler 进行梯度缩放
    scaler.scale(loss).backward()

    # 使用 scaler 进行优化器步进
    scaler.step(optimizer)

    # 更新缩放因子
    scaler.update()

    print(f"Epoch [{epoch+1}/10], Loss: {loss.item()}")

4. 总结

  • torch.cuda.amp.GradScaler 是用于混合精度训练的工具,通过动态缩放损失值来提高数值稳定性。
  • 使用 autocast 上下文管理器来自动处理前向传播中的精度切换。
  • 在反向传播和优化器更新时,通过 scaler 来处理损失缩放和梯度计算。

混合精度训练能够在现代 GPU 上显著提升训练速度和效率,同时通过 GradScaler 保持模型的稳定性和精度。

相关推荐
陈鋆25 分钟前
智慧城市初探与解决方案
人工智能·智慧城市
qdprobot25 分钟前
ESP32桌面天气摆件加文心一言AI大模型对话Mixly图形化编程STEAM创客教育
网络·人工智能·百度·文心一言·arduino
QQ395753323726 分钟前
金融量化交易模型的突破与前景分析
人工智能·金融
QQ395753323727 分钟前
金融量化交易:技术突破与模型优化
人工智能·金融
The_Ticker39 分钟前
CFD平台如何接入实时行情源
java·大数据·数据库·人工智能·算法·区块链·软件工程
Elastic 中国社区官方博客1 小时前
Elasticsearch 开放推理 API 增加了对 IBM watsonx.ai Slate 嵌入模型的支持
大数据·数据库·人工智能·elasticsearch·搜索引擎·ai·全文检索
jwolf21 小时前
摸一下elasticsearch8的AI能力:语义搜索/vector向量搜索案例
人工智能·搜索引擎
有Li1 小时前
跨视角差异-依赖网络用于体积医学图像分割|文献速递-生成式模型与transformer在医学影像中的应用
人工智能·计算机视觉
傻啦嘿哟1 小时前
如何使用 Python 开发一个简单的文本数据转换为 Excel 工具
开发语言·python·excel
B站计算机毕业设计超人1 小时前
计算机毕业设计SparkStreaming+Kafka旅游推荐系统 旅游景点客流量预测 旅游可视化 旅游大数据 Hive数据仓库 机器学习 深度学习
大数据·数据仓库·hadoop·python·kafka·课程设计·数据可视化