PyTorch-混合精度训练

最新推荐文章于 2024-08-31 20:56:25 发布

周先森爱吃素

最新推荐文章于 2024-08-31 20:56:25 发布

阅读量5.7k

点赞数 8

分类专栏： PyTorch 文章标签： PyTorch混合精度训练

本文链接：https://blog.csdn.net/zhouchen1998/article/details/112749953

版权

PyTorch 专栏收录该内容

10 篇文章 12 订阅

订阅专栏

简介

自动混合精度训练（auto Mixed Precision，amp）是深度学习比较流行的一个训练技巧，它可以大幅度降低训练的成本并提高训练的速度，因此在竞赛中受到了较多的关注。此前，比较流行的混合精度训练工具是由NVIDIA开发的A PyTorch Extension(Apex)，它能够以非常简单的API支持自动混合精度训练，不过，PyTorch从1.6版本开始已经内置了amp模块，本文简单介绍其使用。

自动混合精度（AMP）

首先来聊聊自动混合精度的由来。下图是常见的浮点数表示形式，它表示单精度浮点数，在编程语言中的体现是float型，显然从图中不难看出它需要4个byte也就是32bit来进行存储。深度学习的模型数据均采用float32进行表示，这就带来了两个问题：模型size大，对显存要求高；32位计算慢，导致模型训练和推理速度慢。

在这里插入图片描述

那么半精度是什么呢，顾名思义，它只用16位即2byte来进行表示，较小的存储占用以及较快的运算速度可以缓解上面32位浮点数的两个主要问题，因此半精度会带来下面的一些优势：

显存占用更少，模型只有32位的一半存储占用，这也可以使用更大的batch size以适应一些对大批尺寸有需求的结构，如Batch Normalization；
计算速度快，float16的计算吞吐量可以达到float32的2-8倍左右，且随着NVIDIA张量核心的普及，使用半精度计算已经比较成熟，它会是未来深度学习计算的一个重要趋势。

那么，半精度有没有什么问题呢？其实也是有着很致命的问题的，主要是移除错误和舍入误差两个方面，具体可以参考这篇文章，作者解析的很好，我这里就简单复述一下。

溢出错误

FP16的数值表示范围比FP32的表示范围小很多，因此在计算过程中很容易出现上溢出（overflow）和下溢出（underflow）问题，溢出后会出现梯度nan问题，导致模型无法正确更新，严重影响网络的收敛。而且，深度模型训练，由于激活函数的梯度往往比权重的梯度要小，更容易出现的是下溢出问题。

舍入误差

舍入误差（Rounding Error）指的是当梯度过小，小于当前区间内的最小间隔时，该次梯度更新可能会失败。上面说的知乎文章的作者用来一张很形象的图进行解释，具体如下，意思是说在 $2^{-3}$ 到 $2^{-2}$ 之间， $2^{-3}$ 每次变大都会至少加上 $2^{-13}$ ，显然，梯度还在这个间隔内，因此更新是失败的。

在这里插入图片描述

那么这两个问题是如何解决的呢，思路来自于NVIDIA和百度合作的论文，我这里简述一下方法：混合精度训练和损失缩放。前者的思路是在内存中使用FP16做储存和乘法运算以加速计算，用FP32做累加运算以避免舍入误差，这样就缓解了舍入误差的问题；后者则是针对梯度值太小从而下溢出的问题，它的思想是：反向传播前，将损失变化手动增大 $2^k$ 倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出；反向传播后，将权重梯度缩小 $2^k$ 倍，恢复正常值。

研究人员通过引入FP32进行混合精度训练以及通过损失缩放来解决FP16的不足，从而实现了一套混合精度训练的范式，NVIDIA以此为基础设计了Apex包，不过Apex的使用本文就不涉及了，下一节主要关注如何使用torch.cuda.amp实现自动混合精度训练，不过这里还需要补充的一点就是目前混合精度训练支持的N卡只有包含Tensor Core的卡，如2080Ti、Titan、Tesla等。

PyTorch自动混合精度

PyTorch对混合精度的支持始于1.6版本，位于torch.cuda.amp模块下，主要是torch.cuda.amp.autocast和torch.cuda.amp.GradScale两个模块，autocast针对选定的代码块自动选取适合的计算精度，以便在保持模型准确率的情况下最大化改善训练效率；GradScaler通过梯度缩放，以最大程度避免使用FP16进行运算时的梯度下溢。官方给的使用这两个模块进行自动精度训练的示例代码链接给出，我对其示例解析如下，这就是一般的训练框架。

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

# 创建梯度缩放器
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        # 通过自动类型转换进行前向传播
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
        # 缩放大损失，反向传播不建议放到autocast下，它默认和前向采用相同的计算精度
        scaler.scale(loss).backward()
        # 先反缩放梯度，若反缩后梯度不是inf或者nan，则用于权重更新
        scaler.step(optimizer)
        # 更新缩放器
        scaler.update()

下面我以简单的MNIST任务做测试，使用的显卡为RTX 3090，代码如下。该代码段中只包含核心的训练模块，模型的定义和数据集的加载熟悉PyTorch的应该不难自行补充。

model = Model()
model = model.cuda()
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

n_epochs = 30
start = time.time()
for epoch in range(n_epochs):
    total_loss, correct, total = 0.0, 0, 0
    model.train()
    for step, data in enumerate(data_loader_train):
        x_train, y_train = data
        x_train, y_train = x_train.cuda(), y_train.cuda()
        outputs = model(x_train)
        _, pred = torch.max(outputs, 1)
        loss = loss_fn(outputs, y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        total += len(y_train)
        correct += torch.sum(pred == y_train).item()

    print("epoch {} loss {} acc {}".format(epoch, total_loss, correct / total))