混合精度训练笔记

最新推荐文章于 2024-08-27 10:48:12 发布

六等分的商鞅

最新推荐文章于 2024-08-27 10:48:12 发布

阅读量790

点赞数 1

分类专栏： tensorflow 文章标签：深度学习机器学习神经网络

本文链接：https://blog.csdn.net/qq_37326944/article/details/122916327

版权

tensorflow 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

同时使用 16 位和 32 位浮点类型训练，加快运行速度，减少内存使用的一种训练方法。
效率会得到提升的前提条件：
1.在batch足够大，2.gpu满足计算能力至少为 7.0 的 NVIDIA GPU，3.模型足够大。
可用字符串 ‘mixed_float16’ 构造策略，并将其设置为全局策略。这会导致随后创建的层使用 float16 和 float32 的混合精度

`policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

`
损失放大：因为float16比32数据动态范围小，所以可能会出现下溢问题，在正向传递中出现下溢的情形更是十分罕见。但是，在反向传递中，梯度可能因下溢而变为零，所以要用损失放大避免下溢。
损失放大的基本概念非常简单：只需将损失乘以某个大数字即可，如 1024，我们将该数字称为***损失标度***。这会将梯度放大1024 倍，大大降低了发生下溢的几率。计算出最终梯度后，将其除以即可得到正确值。
但是实际情况这个标度很难选，如果损失标度太小，梯度可能仍会因下溢而变为零。如果太大，则会出现相反的问题：梯度可能因溢出而变为无穷大。所以引入了选择损失标度，如果用model.fit则会自动选，而自定义训练就得手动选了。

选择损失标度:mixed_float16 策略的损失标度是 tf.mixed_precision.experimental.DynamicLossScale，它会动态确定损失标度值。其他策略在默认情况下没有损失标度。可以用此代码查询损失标度：

loss_scale = policy.loss_scale
print('Loss scale: %s' % loss_scale)

在构建 dtype 策略时，您也可以通过传递一个数字来使用静态损失标度。

new_policy = mixed_precision.Policy('mixed_float16', loss_scale=1024)
print(new_policy.loss_scale)

自定义训练循环训练模型：
使用 mixed_float16，需要将损失放大。
您将使用 tf.keras.mixed_precision.experimental.LossScaleOptimizer 类，其中会封装一个优化器并应用损失放大。该类有两个参数：优化器和损失标度。
像下面这样构造一个类即可使用动态损失标度

optimizer = keras.optimizers.RMSprop()
optimizer = mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic')

传递 ‘dynamic’ 等效于传递 tf.mixed_precision.experimental.DynamicLossScale()
一般我们在自定义循环中这样定义
为了放大损失和缩小梯度，将使用损失标度优化器的两个新方法：

get_scaled_loss(loss)：将损失值乘以损失标度值
get_unscaled_gradients(gradients)：获取一系列放大的梯度作为输入，并将每一个梯度除以损失标度，从而将其缩小为实际值

为了防止梯度发生下溢，必须使用这些函数。随后，如果全部没有出现 Inf 或 NaN 值，则 LossScaleOptimizer.apply_gradients 会应用这些梯度。它还会更新损失标度，如果梯度出现 Inf 或 NaN 值，则会将其减半，而如果出现零值，则会增大损失标度。

@tf.function
def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
    scaled_loss = optimizer.get_scaled_loss(loss)
  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
  gradients = optimizer.get_unscaled_gradients(scaled_gradients)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

在训练的开始阶段，LossScaleOptimizer 可能会跳过前几个步骤。先使用非常大的损失标度，以便快速确定最佳值。经过几个步骤后，损失标度将稳定下来，这时跳过的步骤将会很少。这一过程是自动执行的，不会影响训练质量。

GPU 上使用混合精度时的一些性能提示：
1.16位比32位数据内存少用一半所以可以加大batch
2.gpu中有个叫tensor核心的硬件，可以非常快速地执行 float16 矩阵乘法运算，但是要求张量的某些维度是 8 的倍数。如：

tf.keras.layers.Dense(units=64)
tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)
    其他卷积层也是如此，如 tf.keras.layers.Conv3d
tf.keras.layers.LSTM(units=64)
    其他 RNN 也是如此，如 tf.keras.layers.GRU
tf.keras.Model.fit(epochs=2, batch_size=128)

3.xla：款可以进一步提高混合精度性能，也可以在较小程度上提高 float32 性能的编译器