混合精度训练

Mixed Precision Training

典型的深度学习库在分配数据类型时非常保守。默认分配的数据类型通常是 64 位或 32 位，而其实也存在使用 16 位的空间。从下面的代码中也可以看出这一点：

因此，我们在高效分配内存方面并非完全最优。当然，这样做是为了保证信息表示的更高精度。

然而，这种精度总是以额外的内存使用为代价，而这在所有情况下未必都是我们想要的。

事实上，还可以观察到，许多张量运算（尤其是矩阵乘法）在更小精度数据类型下比在更大精度下要快得多，如下所示：

此外，由于 float16 只有 float32 的一半大小，使用它可以减少训练网络所需的内存。这也让我们能够训练更大的模型、在更大的小批量上训练（从而获得更多加速）等等。混合精度训练（mixed precision training）是工业界实现这些目标的一种相当可靠且被广泛采用的技术。顾名思义，其思想是在可行之处（如卷积和矩阵乘法）使用更低精度的 float16，并与 float32 搭配使用——这也是“混合精度”得名的原因。下面是一些我发现的、使用混合精度训练的模型列表：

很明显，混合精度训练的使用要普遍得多，只是我们平时不太会听到它。

从上面的讨论中应当清楚，当我们使用低精度数据类型（float16）时，可能会在不知不觉中引入一些数值上的不一致和不精确。

为避免这些问题，接下来我想谈谈混合精度训练的一些最佳实践，并附上代码。

在 PyTorch 中使用混合精度训练，需要对现有的网络训练实现做一些修改。假设这是我们当前的 PyTorch 模型训练实现：

我们在这里引入的第一样东西是一个 scaler 对象，用于对损失值进行缩放：

我们这样做是因为，有时候原始损失值可能低到使我们无法在 float16 下以完整精度计算梯度。这种情况可能不会对模型权重产生任何更新。

将损失缩放到更高的数值范围，可以确保即便是很小的梯度也能对权重更新有所贡献。但这些微小的梯度只有在权重矩阵本身以高精度（即 float32）表示时，才能被纳入权重矩阵。因此，作为一种保守做法，我们倾向于将权重保持在 float32 中。

话虽如此，损失缩放这一步并非完全必要，因为根据我的经验，这些微小更新通常出现在模型训练的末期阶段。因此，可以合理地认为，这些小更新未必会对模型性能产生显著影响。但不要把这当作定论，所以我希望你在使用混合精度训练时对此加以验证。

继续，由于权重（也就是矩阵）以 float32 表示，如果它们一直保持这种形式，我们就无法从 float16 表示中获得加速：

为了利用这些基于 float16 的加速，我们遵循以下步骤：

1. 在前向传播期间，我们制作一份 float16 的权重副本。2. 接下来，我们在 float32 下计算损失值，并将其缩放以使梯度（在 float16 下工作）拥有更多精度。a. 我们之所以在 float16 下计算梯度，是因为与前向传播一样，梯度计算也涉及矩阵乘法。b. 因此，将其保持在 float16 可以带来额外的加速。3. 一旦我们在 float16 下计算完梯度，繁重的矩阵乘法运算就已完成。现在，我们要做的就是更新 float32 中的原始权重矩阵。4. 因此，我们为上述梯度制作一份 float32 副本，移除第 2 步施加的缩放，并更新 float32 权重。5. 完成！

前向传播中的混合精度设置由 autocast 上下文管理器完成：

现在，该处理反向传播了。

• 第 13 行 → scaler.scale(loss).backward()：scaler 对象对损失值进行缩放，并调用 backward() 来计算梯度。
• 第 14 行 → scaler.step(opt)：反缩放梯度并更新权重。
• 第 15 行 → scaler.update()：为下一次迭代更新缩放比例。
• 第 16 行 → opt.zero_grad()：清零梯度。

完成！混合精度训练相对于传统训练的成效，从下图中可以明显看出：

混合精度训练比传统训练快两倍多。