当前位置: 首页 > article >正文

PyTorch 中的混合精度训练方法,从 autocast 到 GradScalar

PyTorch 的混合精度训练主要由两个方法实现:amp.autocastamp.GradScalar。在这两个工具的帮助下,可以实现以 torch.float16 的混合精度训练。当然,这两个方法都是模块化并且通常都会一起调用,但并不一定总是需要一起使用。

参考:

Automatic Mixed Precision package - torch.amp
Automatic Mixed Precision examples
Automatic Mixed Precision recipe

amp.autocast

在 PyTorch 中,autocast 是一个用于自动混合精度训练的上下文管理器。它的主要作用是根据操作的类型自动选择使用 float16(半精度)或 float32(单精度)进行计算,从而提高训练速度并减少内存使用。以下是 autocast 的具体功能和工作原理:

  1. 自动选择精度
    操作类型:在深度学习中,不同的操作对数值精度的要求不同。例如,某些操作(如矩阵乘法)在 float16 下可以安全地执行,而其他操作(如某些归一化或损失计算)可能需要 float32 来保持数值稳定性。
    自动化:使用 autocast 时,PyTorch 会自动判断每个操作的最佳精度,并在适当的情况下使用 float16,在其他情况下使用 float32。这使得开发者不需要手动管理每个操作的精度,从而简化了代码。
  2. 上下文管理器
    使用方式autocast 通常用作上下文管理器,使用 with torch.cuda.amp.autocast(): 语句包裹需要进行混合精度计算的代码块。在这个代码块内,所有的张量操作都会根据 autocast 的规则自动选择精度。
  3. 性能提升
    加速训练:通过使用 float16 进行计算,autocast 可以显著提高训练速度,尤其是在支持半精度计算的 GPU 上(如 NVIDIA 的 Volta 和 Ampere 架构)。减少内存使用:使用 float16 还可以减少显存的占用,使得更大的模型或更大的批量大小成为可能。
  4. 示例
    一般来说,建议在模型前向传递和 loss计算中使用 autocast,而不建议在反向传播和参数更新过程中使用,例如:
# 来源:https://pytorch.org/docs/stable/amp.html#torch.autocast
# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass (model + loss)
    with torch.autocast(device_type="cuda"):
        output = model(input) # 前向传递
        loss = loss_fn(output, target) # loss计算

    # Exits the context manager before backward()
    loss.backward() # 后向传递(计算梯度)
    optimizer.step() # 参数更新
  1. 总结
    autocast 是 PyTorch 中实现混合精度训练的关键工具,它通过自动选择操作的精度来提高训练速度和减少内存使用,同时保持数值稳定性。使用 autocast 可以让开发者更专注于模型的设计和训练,而不必担心每个操作的精度管理。

amp.GradScalar

GradScaler 是 PyTorch 中用于混合精度训练的一个重要组件,主要用于处理梯度的缩放,以确保在使用半精度(float16)进行训练时的数值稳定性。以下是对 GradScaler 的详细介绍:

  1. 背景
    在混合精度训练中,使用 float16 进行计算可以提高速度和减少内存使用,但由于 float16 的数值范围和精度较低,可能会导致梯度下溢(即梯度变得过小而被视为零)或上溢(即梯度变得过大而无法表示)GradScaler 的作用就是解决这个问题。
  2. 主要功能
    动态缩放GradScaler 会在反向传播时动态调整损失值的缩放因子,以避免梯度的数值不稳定。它会根据当前的训练状态自动选择合适的缩放因子。
    防止下溢和上溢:通过将损失值乘以一个缩放因子,GradScaler 可以确保计算出的梯度不会因为数值过小而消失,也不会因为数值过大而溢出。
  3. 使用流程
    使用 GradScaler 的典型流程如下:

初始化:创建 GradScaler 的实例;

scaler = torch.cuda.amp.GradScaler()

前向传播:在前向传播时,使用 autocast 来自动选择精度,不建议使用GradScaler

   with torch.cuda.amp.autocast():
       output = model(inputs)
       loss = compute_loss(output, targets)

反向传播:

  • 使用 scaler.scale(loss).backward() 进行反向传播。这里,scale 方法会将损失值缩放,并计算梯度。
  • 使用 scaler.step(optimizer) 更新模型参数。此步骤会将缩放后的梯度应用到优化器中。
  • 使用 scaler.update() 更新缩放因子,以便在下一次迭代中使用。
  1. 示例代码

以下是一个使用 GradScaler 的简单示例:

import torch
from torch.cuda.amp import GradScaler, autocast

model = ...  # 初始化模型
optimizer = ...  # 初始化优化器
scaler = GradScaler()  # 创建 GradScaler 实例

for data, target in dataloader:
    optimizer.zero_grad()  # 清零梯度

    with autocast():  # 自动选择精度
        output = model(data) # 前向传递
        loss = compute_loss(output, target) # 计算loss

    scaler.scale(loss).backward()  # 缩放损失并反向传播以计算梯度
    scaler.step(optimizer)  # 更新模型参数
    scaler.update()  # 更新缩放因子
  1. 总结
    GradScaler 是 PyTorch 中实现混合精度训练的关键工具,它通过动态缩放损失值来确保在使用 float16 进行训练时的数值稳定性。使用 GradScaler 可以有效地避免梯度下溢和上溢的问题,从而提高训练的可靠性和效率。通过结合 autocastGradScaler,开发者可以在不牺牲模型性能的情况下,充分利用混合精度训练的优势。

http://www.kler.cn/a/580582.html

相关文章:

  • 力扣hot100_二叉树(4)_python版本
  • FastAPI 表单参数与文件上传完全指南:从基础到高级实战 [特殊字符]
  • 如何实现创意的角落切割效果:使用 CSS 和 SVG 创建时尚的网页元素
  • NLP常见任务专题介绍(1)-关系抽取(Relation Extraction, RE)任务训练模板
  • TypeError: Cannot convert object to primitive value
  • JavaScript(最后一个元素的索引就是数组的长度减 1)array.length - 1
  • Java 大视界 -- 基于 Java 的大数据实时数据处理框架性能评测与选型建议(121)
  • 请谈谈 TypeScript 中的接口(interface)和类型别名(type alias),它们的区别是什么?
  • Python | 机器学习中最常用的超参数及使用示例
  • Qt常用控件之树形QTreeWidget
  • 如何使用postman来测试接口
  • Synology 部署的 WordPress 無法升級至最新版本時,可以透過以下改良版指南進行排查和解決。
  • 雪藏HsFreezer(游戏冻结工具) v2.21
  • jsonRPC实现前后端分离
  • delphi 正则提取html中的内容
  • listen EACCES: permission denied 0.0.0.0:811
  • Ubuntu 创建systemd服务
  • LeetCode 2380 二进制字符串重新安排顺序需要的时间
  • Docker Compose国内镜像一键部署dify
  • C#常用的循环语句