使用 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。

3. 对优化器使用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优化后，损失值会逐渐变小。

标签：float16,keras,dtype,Keras,precision,policy,mixed
来源： https://blog.csdn.net/drin201312/article/details/121367482