PyTorch-混合精度训练

简介

自动混合精度训练（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, 0model.train()for step, data in enumerate(data_loader_train):x_train, y_train = datax_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))

我这里采用的是一个很小的模型，又是一个很简单的任务，因此模型都是很快收敛，因此精度上没有什么明显的区别，不过如果是训练大型模型的话，有人已经用实验证明，内置amp和apex库都会有精度下降，不过amp效果更好一些，下降较少。上面的loss变化图也是非常类似的。

再来看存储方面，显存缩减在这个任务中的表现不是特别明显，因为这个任务的参数量不多，前后向过程中的FP16存储节省不明显，而因为引入了一些拷贝之类的，反而使得显存略有上升，实际的任务中，这种开销肯定远小于FP32的开销的。

最后，不妨看一下使用混合精度最关心的速度问题，实际上混合精度确实会带来一些速度上的优势，一些官方的大模型如BERT等训练速度提高了2-3倍，这对于工业界的需求来说，启发还是比较多的。

总结

混合精度计算是未来深度学习发展的重要方向，很受工业界的关注，PyTorch从1.6版本开始默认支持amp，虽然现在还不是特别完善，但以后一定会越来越好，因此熟悉自动混合精度的用法还是有必要的。

超强干货来袭云风专访：近40年码龄，通宵达旦的技术人生

PyTorch-混合精度训练

简介

自动混合精度（AMP）

PyTorch自动混合精度

总结

Published by

风君子

发表回复取消回复

简介

自动混合精度（AMP）

PyTorch自动混合精度

总结

Published by

风君子

发表回复 取消回复

发表回复取消回复