【Pytorch】理解自动混合精度训练

  更大的深度学习模型需要更多的计算能力和内存资源。一些新技术的提出，可以更快地训练深度神经网络。我们可以使用 FP16（半精度浮点数格式）来代替 FP32（全精度浮点数格式），研究人员发现串联使用它们是更好的选择。有的 GPU（例如 Paperspace 提供的 Ampere GPU）甚至可以利用较低级别的精度，例如 INT8。

  混合精度允许半精度训练，同时仍保留大部分单精度网络精度。术语“混合精度技术”是指该方法同时使用单精度和半精度表示。

  在使用 PyTorch 进行自动混合精度 (Amp) 训练的概述中，我们演示了该技术的工作原理，逐步介绍使用 Amp 的过程，并通过代码讨论 Amp 技术的应用。

混合精度概述

  在深度学习的世界里，使用 FP16 进行计算不仅能显著提升性能，还能节省内存。然而，这种方法也带来了两个主要问题：精度溢出和舍入误差。这两个问题是深度学习中 FP16 计算的关键挑战。

   精度溢出（Precision Overflow）：

  在 FP16 格式下，由于位宽较小，可表示的数值范围远小于 FP32 或 FP64。这容易导致数值过大或过小而无法在 FP16 的表示范围内精确表示。在深度学习中，这可能引起梯度消失或梯度爆炸，因为一些小的梯度值可能变成零（下溢），而一些大的梯度值可能变得无限大（上溢）。这种溢出问题会严重影响模型训练的稳定性和最终性能。

  舍入误差（Rounding Error）：

  FP16 由于其16位的表示限制，相比于 FP32 或 FP64，舍入误差更加明显。在深度学习中，每次计算的舍入误差会累积，尤其是在多层和复杂运算中。这可能导致模型输出与使用更高精度计算时存在显著差异。对于那些对精确度要求极高的应用（比如金融或医疗领域），这种误差可能造成不可接受的后果。

  为了缓解这些问题，混合精度训练方法在关键部分（如权重更新）使用 FP32 来保持精度，而在其他操作（如前向传播）中使用 FP16 来提高效率。混合精度训练中，我特别注意到了权重备份（Weight Backup）、损失放大（Loss Scaling）、精度累加（Precision Accumulated）这三种技术的重要性。

  权重备份（Weight Backup）:

  在混合精度训练中，为了确保数值稳定性，模型的权重通常会在 FP16 和 FP32 两种格式下同时维护。权重备份是指保留 FP32 格式的权重副本，这样即使在大部分使用 FP16 格式的计算过程中出现数值不稳定现象，我们仍然能依靠 FP32 权重副本保持稳定和精确。这对于更新模型参数时的准确性至关重要。

  损失放大（Loss Scaling）:

  在混合精度训练中，由于 FP16 的表示范围限制，梯度值可能太小而无法在 FP16 中准确表示，导致有效梯度变为零。损失放大是通过在计算梯度前将损失函数的值乘以一个较大的常数（放大因子），从而放大梯度值，使其在 FP16 范围内可表示且非零。在反向传播后，再将放大的梯度除以相同的放大因子，恢复原始比例，这样可以有效减少梯度下溢问题。

  精度累加（Precision Accumulation）:

  精度累加是指在权重更新过程中，即使梯度计算是在 FP16 下完成的，但权重更新则在 FP32 精度下进行。这有助于减少舍入误差和累积误差，尤其在训练过程中涉及大量累积操作时。由于 FP32 提供更高的数值精度和更大的表示范围，可以更准确地累积小梯度值，避免更新权重时的数值不稳定性。

  综上所述，通过将这些技术相结合，混合精度训练能够有效地利用 FP16 带来的性能优势，同时最大限度地减少精度损失和计算不稳定性。

实验对比

  为了进一步验证这些技术的有效性，我设计了两个实验对比使用混合精度和传统的 FP32 两种方式进行训练。以下是我在这两个实验中使用的代码片段：

FP16与FP32混合训练代码：

import torch
from tensorboardX import SummaryWriter
from torch import optim, nn
import time

from torch.cuda.amp import GradScaler, autocast


class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linears = nn.Sequential(
            nn.Linear(2, 20000),

            nn.Linear(20000, 20000),
            # nn.Dropout(0.1),

            nn.Linear(20000, 200),
            # nn.LayerNorm(20),

            nn.Linear(200, 20),
            # nn.LayerNorm(20),

            nn.Linear(20, 1),
        )

    def forward(self, x):
        _ = self.linears(x)
        return _

lr = 0.0001
iteration = 1000


x1 = torch.arange(-1000, 1000).float().to('cuda')
x2 = torch.arange(0, 2000).float().to('cuda')
x = torch.cat((x1.unsqueeze(1), x2.unsqueeze(1)), dim=1)
y = (2*x1 - x2 + 1).to('cuda')

model = Model().to('cuda')
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.01)
loss_function = torch.nn.MSELoss()

scaler = GradScaler()

start_time = time.time()
writer = SummaryWriter(comment='_FP16')

for iter in range(iteration):
    with autocast():
        y_pred = model(x)
        loss = loss_function(y, y_pred.squeeze())
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    writer.add_scalar('loss', loss, iter)
    optimizer.zero_grad()

    if iter % 100 == 0:
        # 获取 GPU 的内存使用情况
        print("GPU Memory Allocated:", torch.cuda.memory_allocated())
        print("GPU Memory Cached:   ", torch.cuda.memory_reserved())

print("Time: ", time.time() - start_time)
torch.save(model.state_dict(), 'model_state_dict_fp16.pth')

FP32训练代码：

import torch
from tensorboardX import SummaryWriter
from torch import optim, nn
import time

from torch.cuda.amp import GradScaler, autocast


class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linears = nn.Sequential(
            nn.Linear(2, 20000),

            nn.Linear(20000, 20000),
            # nn.Dropout(0.1),

            nn.Linear(20000, 200),
            # nn.LayerNorm(20),

            nn.Linear(200, 20),
            # nn.LayerNorm(20),

            nn.Linear(20, 1),
        )

    def forward(self, x):
        _ = self.linears(x)
        return _

lr = 0.0001
iteration = 1000


x1 = torch.arange(-1000, 1000).float().to('cuda')
x2 = torch.arange(0, 2000).float().to('cuda')
x = torch.cat((x1.unsqueeze(1), x2.unsqueeze(1)), dim=1)
y = (2*x1 - x2 + 1).to('cuda')

model = Model().to('cuda')
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.01)
loss_function = torch.nn.MSELoss()

scaler = GradScaler()

start_time = time.time()
writer = SummaryWriter(comment='_FP32')

for iter in range(iteration):

    y_pred = model(x)
    loss = loss_function(y, y_pred.squeeze())
    loss.backward()
    optimizer.step()


    writer.add_scalar('loss', loss, iter)
    optimizer.zero_grad()

    if iter % 100 == 0:
        # 获取 GPU 的内存使用情况
        print("GPU Memory Allocated:", torch.cuda.memory_allocated())
        print("GPU Memory Cached:   ", torch.cuda.memory_reserved())

print("Time: ", time.time() - start_time)
torch.save(model.state_dict(), 'model_state_dict_fp32.pth')

  最终两者的效果如下：

实验名	占用GPU	消耗时间
FP16与FP32混合训练	4867271680 bytes	78.73 s
FP32训练	4867274752 bytes	140.18 s

  实验分析如下：

  1、在内存占用方面，两种训练方法几乎相同。这可能是因为模型结构和数据集大小相同，所以内存占用没有显著差异。然而，通常情况下，FP16训练应该占用更少的内存，因为它使用的是半精度浮点数。

  2、在训练时间方面，混合精度训练明显快于纯FP32训练。这是因为FP16训练可以加快计算速度并降低内存需求，从而允许模型更快地运行。混合精度训练结合了FP16的高效率和FP32的数值稳定性，提供了一个平衡的解决方案。

  在提供的损失图中，我们可以看到蓝色曲线代表使用全精度（FP32）训练的模型损失，而橙色曲线代表使用混合精度（FP16与FP32）训练的模型损失。以下是对两种训练方法损失曲线的进一步分析：

  · 在下降到一定程度后，两条曲线都达到平稳状态，这表明模型已基本收敛。在这个平稳阶段，损失变化不大，说明模型在训练集上的表现已经稳定。

  · 没有明显的过拟合迹象，因为损失曲线没有再次上升的趋势。

  · 混合精度训练似乎在时间效率略优于全精度训练，尽管最终损失值的差异不大，但在追求快速迭代和高效训练的情况下，选择混合精度训练会更有优势；而在收敛速度上，混合精度训练在140epoch时收敛完毕，而全精度训练早在100epoch就收敛完，说明全精度训练收敛较于混合精度训练，在step层面更快，而它在时间层面更慢。