可以使用混合精度 mixed precision 给 Keras 加速,3个操作步骤如下:
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。
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点: