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

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

1. 使用算力在 7.0以上的GPU，比如 NVIDIA 的 RTX 3090 等。
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 。 代码如下 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. 不同的模型和硬件配置，加速程度也会不一样。我在一次模型试验中，看到了 1.25倍的加速，相当于5天的计算量，在使用混合精度后，4天可以算完。并且这里使用的 Keras 版本是 2.4版，即 TensorFlow 2.4 。
2. 损失值在计算过程中会出现个别 inf 或 nan 的情况，这是因为某些 epoch 中损失值过大引起的，在模型经过若干 epochs 优化后，损失值会逐渐变小。