使用混合精度 mixed precision 给 Keras 加速

最新推荐文章于 2024-08-29 10:09:00 发布

杭州的平湖秋月

最新推荐文章于 2024-08-29 10:09:00 发布

阅读量2.7k

点赞数

CC 4.0 BY-SA版权

分类专栏：深度学习文章标签：深度学习 tensorflow keras python

本文链接：https://blog.csdn.net/drin201312/article/details/121367482

深度学习专栏收录该内容

6 篇文章

订阅专栏

通过设置全局策略为mixed_float16并配合LossScaleOptimizer，可以在支持的GPU上加速Keras模型的训练。混合精度计算利用float16提高速度，同时保持float32的数值稳定性。在实验中，模型训练速度提升约1.25倍，减少了计算时间。需要注意的是，可能会遇到损失值溢出导致的inf或nan，这通常在训练后期会改善。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

可以使用混合精度 mixed precision 给 Keras 加速，3个操作步骤如下：

使用算力在 7.0以上的GPU，比如 NVIDIA 的 RTX 3090 等。
在建立模型之前，设置 global_policy 为 mixed_float16。代码如下。

policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)
# 另一种写法是只用下面这一行语句。但是为了验证 policy.compute_dtype，必须使用上面2行语句。
# tf.keras.mixed_precision.set_global_policy("mixed_float16")

# 上面设置好 global_policy 之后，可以用下面2行验证数据类型。
print(f'Compute dtype: {policy.compute_dtype}')
print(f'Variable dtype: {policy.variable_dtype}')
# 变量的数据类型保持为 float32，以保证数字的稳定性 numeric stability。

对优化器使用 LossScaleOptimizer 。代码如下 2 行。

opt_adam = keras.optimizers.Adam(learning_rate=lr, beta_1=0.9, beta_2=0.999)                
opt_mixed_precision = keras.mixed_precision.LossScaleOptimizer(opt_adam)
# 后续编译模型时要使用 opt_mixed_precision，即 model.compile(optimizer=opt_mixed_precision ……

混合精度的原理，简单来说，就是在计算时从 float32 格式改为使用 float16 格式，从而加快计算速度（而保存的数据本身依然是使用 float32 格式）。
此外，为了不丢失过大或过小的损失值，避免损失值溢出，还需要使用 LossScaleOptimizer 对损失值进行动态缩小和放大，使得损失值能够用 float16 格式进行计算。
对混合精度的详细介绍，可以参看官方文档 https://keras.io/api/mixed_precision 和 https://tensorflow.google.cn/guide/mixed_precision

最后注意2点：

不同的模型和硬件配置，加速程度也会不一样。我在一次模型试验中，看到了 1.25倍的加速，相当于5天的计算量，在使用混合精度后，4天可以算完。并且这里使用的 Keras 版本是 2.4版，即 TensorFlow 2.4 。
损失值在计算过程中会出现个别 inf 或 nan 的情况，这是因为某些 epoch 中损失值过大引起的，在模型经过若干 epochs 优化后，损失值会逐渐变小。