启用mixed_float16混合精度时TensorFlow Hub示例运行报错

参考示例代码加载TensorFlow Hub中的模型时，FP32精度模式下代码运行完全正常。添加tf.keras.mixed_precision.set_global_policy('mixed_float16')启用Float16混合精度后，程序抛出错误。报错初看疑似维度不匹配，但相同代码在FP32模式下运行完全正常，无维度相关问题。

复现代码 #

import tensorflow as tf
import tensorflow_hub as hub
IMAGE_SIZE = (224,224)

class_names = ['cat','dog']

# 注释掉以下行时，代码可正常运行
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# --------

model_handle = "https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/5"
do_fine_tuning = False
print("Building model with", model_handle)
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    hub.KerasLayer(model_handle, trainable=do_fine_tuning),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(len(class_names),
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))
])
model.build((None,)+IMAGE_SIZE+(3,))
model.summary()

核心报错 #

运行代码抛出的核心错误为类型不匹配：传入KerasLayer的输入张量为float16类型，但SavedModel加载的concrete function仅支持float32类型输入，无法找到匹配的函数签名。

ValueError: Could not find matching concrete function to call loaded from the SavedModel. Got:
  Positional arguments (4 total):
    * <tf.Tensor 'inputs:0' shape=(None, 224, 224, 3) dtype=float16>
    * False
    * False
    * 0.99
  Keyword arguments: {}

 Expected these arguments to match one of the following 4 option(s):
（所有可选签名均要求第一个输入张量为dtype=tf.float32类型）

该问题的本质是TensorFlow Hub上多数旧版预训练SavedModel的对外输入签名仅支持float32，全局开启混合精度后，Keras会自动将上游层输出转为float16，和预训练层的输入要求冲突。只需在hub.KerasLayer前增加一层显式类型转换，将输入转为float32即可，该操作不会抵消混合精度的加速效果：预训练层内部支持低精度计算的算子仍会自动使用float16运算，仅层入口处做一次类型转换，开销可以忽略。
修改后的模型构建代码如下：

model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    # 新增类型转换层，适配TF Hub模型的float32输入要求
    tf.keras.layers.Lambda(lambda x: tf.cast(x, tf.float32)),
    hub.KerasLayer(model_handle, trainable=do_fine_tuning),
    tf.keras.layers.Dropout(rate=0.2),
    # 混合精度训练建议最终分类层显式指定float32类型，保证数值稳定性
    tf.keras.layers.Dense(len(class_names),
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001),
                          dtype=tf.float32)
])

注意事项 #

• 无需关闭全局混合精度策略，该修改仅调整预训练层的入口输入类型，其余算子仍按混合精度策略执行，可正常获得训练/推理加速收益
• 若对预训练模型做微调，该写法同样适用，不会影响反向传播的精度适配
• 若使用新版原生支持float16输入的TF Hub模型（通常模型卡会明确标注支持mixed precision），可去掉该类型转换层进一步降低开销

内容的提问来源于stack exchange，提问作者learner