使用 mixed precision 给 Keras 加速

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

  1. 使用算力在 7.0以上的GPU,比如 NVIDIA的 RTX 3090, 3080等。
  2. 在建立模型之前,设置 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。
  1. 对优化器使用LossScaleOptimizer 。
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)

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

最后注意2点:

  1. 不同的模型,加速程度也会不一样。我在一次模型在实验中,能够实现加速 1.25倍,相当于5天的计算量,在使用混合精度后,4天可以算完。
  2. 损失值在计算过程中会出现个别 inf或nan的情况,这是因为某些epoch中损失值过大引起的,在模型经过若干epochs优化后,损失值会逐渐变小。
上一篇:BugKu之备份是个好习惯


下一篇:git reset