使用PyTorch实现自动混合精度

教程

PyTorch

简介

更大的深度学习模型需要更多的计算能力和内存资源。通过新技法的开发，深度神经网络的训练速度得到了提升。而不是使用FP32（全精度浮点数格式），你可以使用FP16（半精度浮点数格式），研究者们发现这两种格式结合使用效果更佳。

混合精度允许在保持单精度网络准确度的同时进行半精度训练。术语“混合精度技术”是指这种方法同时利用了单精度和半精度表示。

在本文中，我们将概述使用PyTorch进行的自动混合精度（AMP）训练，展示这一技术是如何工作的，逐步讲解如何使用AMP，并讨论AMP技术的更多高级应用，为用户提供代码支架以便他们后来整合到自己的代码中。

先决条件

PyTorch基础知识：熟悉PyTorch，包括其核心概念如张量、模块和训练循环。

深度学习基础知识理解：如神经网络、反向传播和优化的概念。

混合精度训练的知识：了解混合精度训练的优势和劣势，包括减少内存使用和加快计算速度。

访问兼容硬件：一款支持混合精度的GPU，例如具有张量核心（如Volta、Turing、Ampere架构）的NVIDIA GPU。

Python和CUDA设置：安装有PyTorch并且配置了CUDA以实现GPU加速的工作Python环境。

混合精度概述

与大多数深度学习框架一样，PyTorch通常在32位浮点数（FP32）上进行训练。然而，FP32并非总是成功的必要条件。可以对少数操作使用16位浮点数，在这些操作中，FP32消耗更多的时间和内存。

因此，NVIDIA工程师开发了一种技术，允许在FP32中对少数操作进行混合精度训练，而网络的大部分运行在FP16中。

尽可能地将模型转换为使用float16数据类型。
保持float32主权重，以在每个迭代中累积权重更新。
使用损失缩放来保持微小梯度值。

对于混合精度训练，PyTorch已经内置了许多功能。
当你调用模块的.half()方法时，模块的参数会被转换为FP16；当你对张量调用.half()方法时，张量的数据也会被转换为FP16。混合精度算术将用于执行这些模块或张量上的任何操作。PyTorch很好地支持了NVIDIA数学库（cuBLAS和cuDNN）。FP16管道中的数据使用Tensor Cores来执行GEMM和卷积。要在cuBLAS中使用Tensor Cores，GEMM的尺寸([M, K] x [K, N] -> [M, N])必须是8的倍数。

介绍Apex

Apex的混合精度工具旨在保持单精度训练的准确性和稳定性的同时，提高训练速度。Apex可以在FP16或FP32下执行操作，自动处理主参数转换，并自动缩放损失。

Apex 是为了让研究人员更容易地将混合精度训练包含在他们的模型中而创建的。Amp（自动混合精度）是 Apex 的一个特性，这是一个轻量级的 PyTorch 扩展。用户只需在他们的网络中添加几行代码，就可以利用 Amp 享受混合精度训练的好处。Apex 是在 CVPR 2018 上推出的，值得注意的是，自从发布以来，PyTorch 社区对 Apex 表现出了强烈的支持。

通过对运行模型进行微小更改，Amp 使得在创建或运行脚本时无需担心混合类型。Amp 的假设可能不适合那些以不寻常方式使用 PyTorch 的模型，但提供了相应的挂钩来调整这些假设。

Amp 提供了混合精度训练的所有优势，而无需显式管理损失缩放或类型转换。Apex 的 GitHub 网站提供了安装说明，其官方 API 文档可以在此处找到。

Amp 是如何工作的

Amp在逻辑层面上使用白名单/黑名单范式。PyTorch中的张量操作包括神经网络函数，如torch.nn.functional.conv2d，简单的数学函数，如torch.log，以及张量方法，如torch.Tensor.add__。在这个体系中有三类主要的函数:

白名单: 可能受益于FP16数学速度提升的函数。典型应用包括矩阵乘法和卷积。
黑名单: 对于可能不够16位精度的函数，输入应为FP32。
其他所有（剩下的任何函数）： 可以在FP16下运行的函数，但将它们在FP32 -> FP16之间转换以执行的成本不值得，因为提速不显著。

Amp的任务在理论上是简单的。Amp在调用PyTorch函数之前，确定其是否在白名单、黑名单或两者皆不是。如果在白名单上，所有参数应转换为FP16；如果在黑名单上，所有参数应转换为FP32。如果都不是，只需确保所有参数类型一致。这个策略在实际应用中并不像看起来那么简单。

使用Amp与PyTorch模型一起使用

要在当前的PyTorch脚本中包含Amp，请按照以下步骤操作：

使用Apex库导入Amp。
初始化Amp，以便它可以对模型、优化器和PyTorch内部函数进行所需的更改。
注意反向传播（.backward()）发生的位置，以便Amp可以同时缩放损失并清除每迭代状态。

第一步

第一步只需一行代码：

from apex import amp

第二步

用于训练的神经网络模型和优化器必须已经指定以完成这一步，这一步只有一行代码。

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

附加设置允许您对Amp的张量和操作类型调整进行微调。函数amp.initialize()接受许多参数，我们只需指定其中的三个：

models (torch.nn.Module 或torch.nn.Module的列表) – 要修改/转换的模型。
优化器（可选，torch.optim.Optimizer 或 torch.optim.Optimizers 的列表）– 要修改/转换的优化器。训练时必填，推理时可选。
opt_level（字符串，可选，默认=“O1”）– 纯或混合精度优化级别。接受的值有“O0”，“O1”，“O2”和“O3”，以上有详细解释。有四个优化级别：

O0 针对 FP32 训练： 这是一个无操作。由于您的传入模型应该是 FP32 already，O0 可能有助于建立准确性的基线，所以无需担心这个问题。

O1 针对混合精度（建议用于典型场景）： 修改所有 Tensor 和 Torch 方法，以使用白名单-黑名单输入转换方案。在 FP16 中，白名单操作（如 Tensor Core 友好的操作，如 GEMMs 和卷积）被执行。例如，softmax 是一个黑名单操作，需要 FP32 精度。除非有明确的说明，否则 O1 还采用动态损失缩放。

O2 针对“几乎 FP16”混合精度： O2 将模型权重转换为 FP16，修补模型的前向方法以将输入数据转换为 FP16，保持批归一化为 FP32，保持 FP32 主权重，更新优化器的 param_groups，使 optimizer.step() 直接对 FP32 权重进行操作，并实现动态损失缩放（除非被覆盖）。与 O1 不同，O2 不修补 Torch 函数或 Tensor 方法。

O3用于FP16训练：O3在真正的混合精度方面可能不如O1和O2稳定。因此，为您的模型设置一个基线速度是有益的，以便可以评估O1和O2的效率。
O3中额外的属性覆盖keep_batchnorm_fp32=True可能会帮助您确定如果您的模型使用批量归一化，“速度 of light”，从而启用cudnn批量归一化。

O0和O3都不是真正的混合精度，但它们分别有助于设置精度和速度基线。混合精度的实现被定义为O1和O2。
您可以尝试两者，看看哪一种最有利于您特定模型的性能和精度提升。

步骤3

确保您在代码中标识出反向传播发生的位置。
会出现几行类似于这样的代码：

loss = criterion(…)
loss.backward()
optimizer.step()

步骤4

使用Amp上下文管理器，您可以简单地通过包装反向传播来启用损失缩放：

loss = criterion(…)
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

就是这样。现在您可以启用混合精度训练重新运行您的脚本。

捕获函数调用

由于PyTorch非常灵活和动态，它缺乏静态模型对象或图来捕获并插入上述提到的转换。通过“monkey patching”所需函数，Amp可以动态地拦截和转换参数。

例如，您可以使用下面的代码来确保torch.nn.functional.linear方法中的参数总是转换为fp16:

orig_linear = torch.nn.functional.linear
def wrapped_linear(*args):
 casted_args = []
  for arg in args:
    if torch.is_tensor(arg) and torch.is_floating_point(arg):
      casted_args.append(torch.cast(arg, torch.float16))
    else:
      casted_args.append(arg)
  return orig_linear(*casted_args)
torch.nn.functional.linear = wrapped_linear

尽管Amp可能添加了一些优化，以使代码更具弹性，但调用Amp.init()实际上会导致所有相关PyTorch函数中插入monkey patches，以便在运行时正确转换参数。

最小化转换

由于Amp保留所有参数转换的内部缓存并根据需要重新使用它们，因此每个权重在每个迭代中只转换一次FP32 -> FP16。在每个迭代中，反向传播的上下文管理器告诉Amp何时清除缓存。

使用PyTorch进行自动转换和梯度缩放

“自动混合精度训练”是指将torch.cuda.amp.autocast和torch.cuda.amp.GradScaler结合起来使用。使用torch.cuda.amp.autocast，您可以只为特定的部分设置自动转换。自动转换会自动为GPU操作选择精度，以优化效率的同时保持准确性。

torch.cuda.amp.GradScaler实例使得执行梯度缩放步骤更加容易。梯度缩放可以减少梯度下溢，有助于使用float16梯度的网络实现更好的收敛。

以下是一些演示如何在PyTorch中使用autocast()获得自动混合精度的代码：

 # 用默认精度创建模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

 # 在训练开始时创建一个GradScaler。
scaler = GradScaler()

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

         # 使用autocast执行前向传播。
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

         # 后向操作使用autocast为相应的前向操作选择的相同数据类型。
        scaler.scale(loss).backward()

         # 首先调用scaler.step()对优化器分配的参数的梯度进行去缩放。
   
        scaler.step(optimizer)

         # 更新下一轮迭代的缩放比例。
        scaler.update()

如果某个特定操作的前向传播具有float16输入，那么这个操作的反向传播会产生float16梯度，而float16可能无法表示小梯度的 magnitude。

如果这些值被冲刷为零（“下溢”），则会丢失相关参数的更新。

梯度缩放是一种使用缩放因子的技术，用以乘以网络的损失，然后在对缩放后的损失进行反向传播以避免下溢。它还需要通过相同的因子缩放网络中流动的反向梯度。因此，梯度值具有更大的 magnitude，从而防止它们被冲刷为零。

在更新参数之前，应该将每个参数的梯度（.grad属性）进行缩放，以防止缩放因子干扰学习率。由于它们是模块化的，autocast和GradScaler都可以独立使用。

与未缩放的梯度一起工作

梯度裁剪

我们可以通过使用`Scaler.scale(Loss).backward()`方法来缩放所有梯度。在`backward()`和`scaler.step(optimizer)`之间的参数的`.grad`属性在更改或检查它们之前必须进行解缩放。如果您想要限制梯度集的全局范数（请参阅torch.nn.utils.clip_grad_norm_()）或最大幅度（请参阅torch.nn.utils.clip_grad_value_()）小于或等于某个值（某些用户设定的阈值），您可以使用一种称为“梯度裁剪”的技术。

不进行解缩放的裁剪会导致梯度的范数/最大幅度被缩放，从而使您请求的阈值无效（该阈值应该是未缩放的梯度的阈值）。优化器给定的参数的梯度通过`scaler.unscale (optimizer)`进行解缩放。
您可以使用scaler.unscale (optimizer1)来解缩放之前传递给另一个优化器（如optimizer1）的其他参数的梯度。我们可以通过添加两行代码来说明这个概念：

# 原地解缩放优化器分配的参数的梯度
        scaler.unscale_(optimizer)
# 由于优化器分配的参数的梯度已经解缩放，因此像往常一样剪辑： 
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

使用缩放梯度工作

梯度累积

梯度累积基于一个极其基础的概念。它不是更新模型参数，而是等待并累积连续批次的梯度来计算损失和梯度。

在一定批次数之后，根据累积梯度来更新参数。以下是使用梯度累积的pytorch的代码片段：

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            # 标准化损失 
            loss = loss / iters_to_accumulate

        # 累积缩放梯度。
        scaler.scale(loss).backward()
          # 权重更新
        if (i + 1) % iters_to_accumulate == 0:
            # 如果需要，可以在這裡取消缩放 
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

梯度累积在batch_per_iter * iters_to_accumulate的适当批量大小的批次之间累积梯度。
缩放应该针对有效批次进行校准；这包括检查inf/NaN梯度，如果检测到任何inf/NaN，则跳过步骤，并将缩放更新到有效批次的粒度。
在将特定有效批次的梯度相加时，保持梯度以缩放和一致的缩放因子也很重要。

如果在累积完成之前梯度未进行缩放（或者缩放因子发生变化），下一次反向传播将会将缩放的梯度添加到未缩放的梯度（或者缩放因子不同的梯度）中，此时无法恢复累积的未缩放梯度。必须在应用累积的未缩放梯度步骤前进行。

你可以在每步之前不久使用unscale函数对梯度进行去缩放。unscale之后，所有即将到来的步的缩放梯度都被累积了。
为了确保完整的有效批次，你只需在每个迭代的末尾调用update，而不是之前调用的step。
enumerate(data)函数允许我们在遍历数据时跟踪批次索引。
将运行损失除以iters_to_accumulate(loss / iters_to_accumulate)。这通过规范化损失来减少我们处理的每个小批次的贡献。如果你在每个批次内平均损失，除法已经是正确的，不需要进一步规范化。这个步骤取决于你是如何计算损失的，可能是不必要的。
当我们使用 `scaler.scale(loss).backward()` 时，PyTorch 会累积缩放的梯度并将其存储起来，直到我们调用 `optimizer.zero_grad()`。

梯度惩罚

在实现梯度惩罚时，使用 `torch.autograd.grad()` 来构建梯度，将这些梯度组合形成惩罚值，然后将其添加到损失中。下面是一个没有缩放或自动 casting 的 L2 惩罚的例子。

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)

        # 创建梯度
        grad_prams = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # 计算惩罚项并将其添加到损失中
        grad_norm = 0
        for grad in grad_prams:
            grad_norm += grad.pow(2).sum()
        grad_norm = grad_norm.sqrt()
        loss = loss + grad_norm

        loss.backward()

        # 您可以在這裡裁剪梯度

        optimizer.step()

提供给 `torch.autograd.grad()` 的张量应该缩放到实施梯度惩罚。在组合它们以获得惩罚值之前，需要反缩放梯度。由于惩罚项的计算是前向传播的一部分，它应该在自动 cast 上下文中进行。
对于相同的 L2 惩罚，以下是它的外观：

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # 为 autograd.grad 的反向传播执行损失缩放，结果 #scaled_grad_prams
        scaled_grad_prams = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        # 在计算惩罚之前创建 grad_prams（grad_prams 必须是 #未缩放的）。 
        # 因为没有任何优化器拥有 scaled_grad_prams，所以使用传统的除法而不是 scaler.unscale_：
        inv_scaled = 1./scaler.get_scale()
        grad_prams = [p * inv_scaled for p in scaled_grad_prams]

        # 计算并添加惩罚项到损失中。 
        with autocast():
            grad_norm = 0
            for grad in grad_prams:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        # 对反向调用应用缩放。
        # 适当地缩放叶节点梯度。
        scaler.scale(loss).backward()

        # 您可以在這裡进行 unscale_ 

        # step() 和 update() 像往常一样进行。
        scaler.step(optimizer)
        scaler.update()

使用多个模型、损失和优化器工作

如果你的网络中有许多损失，那么必须在每个损失上调用Scaler.scale。
如果你的网络中有许多优化器，你可以在它们中的任何一个上调用scaler.unscale，并且必须在每个优化器上调用scaler.step。然而，scaler.update应该只使用一次，在所有优化器完成本迭代步骤之后：

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        with autocast():
            output1 = model1(input)
            output2 = model2(input)
            loss1 = loss_fn(2 * output1 + 3 * output2, target)
            loss2 = loss_fn(3 * output1 - 5 * output2, target)

       #尽管保留图与amp无关，但它出现在这个示例中，因为两个backward()调用共享某些图表区域。 
        scaler.scale(loss1).backward(retain_graph=True)
        scaler.scale(loss2).backward()

        #如果你想查看或调整他们拥有的梯度，你可以指定哪些优化器进行显式的去缩放。.
        scaler.unscale_(optimizer1)

        scaler.step(optimizer1)
        scaler.step(optimizer2)

        scaler.update()

每个优化器检查自己的梯度是否存在inf/NaN，并单独决定是否跳过步骤。一些优化器可能会跳过步骤，而其他的可能不会。每跳过几个迭代，就会发生一次跳过步骤，因此它不应该影响收敛。对于具有多个优化器的模型，如果你在添加梯度缩放后看到收敛性较差，可以报告这个问题。

使用多个GPU工作

深度学习模型面临的一个最显著问题是因为它们变得太大，无法在单个GPU上进行训练。在单个GPU上训练一个模型可能需要太长时间，因此需要多GPU训练以便尽可能快地准备模型。一位知名研究人员能够将ImageNet训练周期从两周缩短到18分钟，或者在两周内训练最广泛和最先进的Transformer-XL，而不是四年的时间。

数据并行和分布式数据并行

不牺牲质量的情况下，PyTorch提供了使用方便和控制能力最佳结合。nn.DataParallel和nn.parallel.DistributedDataParallel是PyTorch用于在多个GPU上分布训练的两个特性。您可以使用这些易于使用的包装器和对训练网络进行多GPU更改。

单进程中的DataParallel

在单台机器上，DataParallel有助于将训练扩展到多个GPU。
让我们更仔细地看看DataParallel在实际中是如何工作的。
在利用DataParallel训练神经网络时，会发生以下阶段：

来源：

将小批量数据划分在GPU:0上。
将小批量数据分割并分发到所有可用的GPU。
将模型复制到GPU。
在所有GPU上进行前向传播。
在GPU:0上计算与网络输出相关的损失，并将损失返回给各个GPU。应在每个GPU上计算梯度。
在GPU:0上累加梯度，并应用优化器更新模型。

需要注意的是，这里讨论的关注点仅适用于autocast。GradScaler的行为保持不变。无论torch.nn.DataParallel是否为每个设备创建线程进行前向传播，都不重要。autocast状态在每一个中进行通信，以下将正常工作：

model = Model_m()
p_model = nn.DataParallel(model)

# 在主线程中设置autocast
with autocast():
    # p_model中将有autocast 
    output = p_model(input)
    # loss_fn也将autocast
    loss = loss_fn(output)

每个进程一个GPU的DistributedDataParallel

torch.nn.parallel.DistributedDataParallel的文档建议每个进程使用一个GPU以获得最佳性能。在这种情况下，DistributedDataParallel不会内部启动线程；因此，使用autocast和GradScaler不会受到影响。

DistributedDataParallel，每个进程多个GPU

在此，torch.nn.parallel.DistributedDataParallel可能会启动一个旁线程，在每个设备上运行前向传播，类似于torch.nn.DataParallel。解决方法是一样的：将autocast作为模型前向方法的一部分应用，以确保在旁线程中启用。

结论

在这篇文章中，我们介绍了：

介绍了Apex。
了解了Amps的工作原理。
了解了如何执行梯度缩放、梯度裁剪、梯度累积和梯度惩罚。
了解了如何处理多个模型、损失函数和优化器。
了解了如何在单个进程中使用多个GPU进行DataParallel。

参考资料

https://developer.nvidia.com/blog/apex-pytorch-easy-mixed-precision-training/
https://nvidia.github.io/apex/amp.html
https://discuss.pytorch.org/t/accumulating-gradients/30020
https://towardsdatascience.com/how-to-scale-training-on-multiple-gpus-dae1041f49d2

Source:
https://www.digitalocean.com/community/tutorials/automatic-mixed-precision-using-pytorch