PyTorch学习（4）：混合精度训练（torch.cuda.amp）

1. 前述

混合精度训练（Mixed Precision Training）是一种深度学习训练技术，它使用不同精度的数据类型（如单精度FP32和半精度FP16）来执行训练过程中的不同部分，以达到加速训练、减少内存使用和降低计算开销的目的。这种训练方式可以在保持模型精度的同时，显著提高训练效率。

在混合精度训练中，通常将模型的权重和激活值等参数保存在较高精度的数据类型中（如FP32），以确保模型的准确性和稳定性。而在进行前向传播、反向传播和梯度更新等计算密集型操作时，则使用较低精度的数据类型（如FP16）来执行，以利用硬件加速和减少内存带宽的需求。

PyTorch训练框架为我们提供了一个混合训练的模块，称作AMP，全称为 Automatic Mixed Precision。这个模块旨在帮助开发者更容易地在他们的模型训练过程中使用混合精度训练。 混合精度训练是一种减少模型训练过程中浮点运算精度以节省内存和加速训练的技术。它结合了 16 位和 32 位浮点数（分别为 float16 和 float32）的计算，以此来在保持模型性能的同时提高计算效率和减少内存占用。 使用 torch.cuda.amp 模块可以自动地将模型的一部分转换为低精度表示，同时保持关键部分（如模型权重更新）的高精度计算。这样可以在不牺牲模型性能的前提下，显著提高训练速度和减少显存使用。

2. PyTorch混合精度训练模块（torch.cuda.amp）

torch.cuda.amp为我们提供了一系列的API，帮助我们实现混合精度的训练。

（1）GradScaler

创建梯度缩放器

（2）torch.autocast或amp.autocast

作为上下文管理器，可以确保在其作用域范围内在混合精度下执行运算。这一操作时自动实现的。

（3）scale

将混合精度计算获得的Loss恢复到float32，并计算梯度。

（4）step

调用optimizer的step，优化学习率。

（5）update

更新权重参数。

3. torch.cuda.amp例程

from torch.cuda import amp  # 混合精度加速训练

scaler = amp.GradScaler(enabled=cuda)

Optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum

scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[round(epochs * x) for x in [0.7, 0.9]], gamma=0.1)

scheduler.last_epoch = start_epoch - 1


for epoch in range(start_epoch, epochs):

optimizer.zero_grad()

for i, (imgs, targets, paths, _) in pbar:

    with amp.autocast(enabled=cuda):

        # 前向计算

        pred = model(imgs)

        # 计算Loss

        loss = compute_loss(pred, targets.to(device), model)

    # 缩放Loss, 反向传播

    scaler.scale(loss).backward()

    # Optimize

    # 将梯度值缩放回原尺度后，优化器进行一步优化

    # 内部调用optimizer.step()，模型才会更新

    # 可选: nn.utils.clip_grad_norm_(model.parameters(), 0.1)# [SAI-KEY] 梯度裁剪

    scaler.step(optimizer)

    # 更新scalar的缩放信息

    scaler.update()

    optimizer.zero_grad()

4. NVIDIA开源Apex库混合精度训练（apex.amp）

apex.amp是NVIDIA 的一个开源项目 Apex 中的一个模块，它提供了自动混合精度（Automatic Mixed Precision，简称 AMP）的功能。Apex 是一个 PyTorch 的扩展库，它旨在提供高效的 GPU 运算和模型优化工具，而 AMP 是其中的一个重要特性，用于加速深度学习模型的训练过程。

当然，我们现在用的比较多的，还是前面讲的torch.cuda.amp，此处仅提供一个例程加以说明。

from apex import amp  # 提供了可以使用混合精度的方便方法，以加速训练

Optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum

scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[round(epochs * x) for x in [0.7, 0.9]], gamma=0.1)

scheduler.last_epoch = start_epoch - 1


model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)

for epoch in range(start_epoch, epochs):

    optimizer.zero_grad()

    for i, (imgs, targets, paths, _) in pbar:

    # 前向计算

    pred = model(imgs)

    # 计算Loss

    loss = compute_loss(pred, targets.to(device), model)  

    # 缩放Loss, 反向传播

    with amp.scale_loss(loss, optimizer) as scaled_loss:

        scaled_loss.backward()

    # Optimize

    # 可选: nn.utils.clip_grad_norm_(model.parameters(), 0.1)# [SAI-KEY] 梯度裁剪

    optimizer.step() # 模型会更新

    optimizer.zero_grad()