启用TensorFlow混合精度训练遇dtype不匹配错误，求解决方案

问题描述 #

启用混合精度训练（执行mixed_precision.set_global_policy(policy="mixed_float16")）后，实例化EfficientNetB0模型触发TypeError，提示x和y必须具有相同dtype（检测到tf.float16 != tf.float32）。错误发生在模型内部的x = x / tf.math.sqrt(IMAGENET_STDDEV_RGB)步骤：输入张量为float16类型，但tf.math.sqrt(IMAGENET_STDDEV_RGB)的计算结果为float32类型。切换回float32全局策略可消除错误，但需要保留混合精度训练。

报错栈信息 #

TypeError                                 Traceback (most recent call last)  in 
5 #mixed_precision.set_global_policy(policy="float32")
6 input_shape = (224, 224, 3)
----> 7 base_model = tf.keras.applications.EfficientNetB0(include_top=False)
8 base_model.trainable = False # freeze base model layers
9

4 frames
/usr/local/lib/python3.7/dist-packages/keras/applications/efficientnet.py in EfficientNetB0(include_top, weights, input_tensor, input_shape, pooling, classes, classifier_activation, **kwargs)
559       classes=classes,
560       classifier_activation=classifier_activation,
---> 561       **kwargs)
562
563

/usr/local/lib/python3.7/dist-packages/keras/applications/efficientnet.py in EfficientNet(width_coefficient, depth_coefficient, default_size, dropout_rate, drop_connect_rate, depth_divisor, activation, blocks_args, model_name, include_top, weights, input_tensor, input_shape, pooling, classes, classifier_activation)
332     # original implementation.
333     # See https://github.com/tensorflow/tensorflow/issues/49930 for more details
---> 334     x = x / tf.math.sqrt(IMAGENET_STDDEV_RGB)
335
336   x = layers.ZeroPadding2D(

/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
151     except Exception as e:
152       filtered_tb = _process_traceback_frames(e.__traceback__)
---> 153       raise e.with_traceback(filtered_tb) from None
154     finally:
155       del filtered_tb

/usr/local/lib/python3.7/dist-packages/keras/layers/core/tf_op_layer.py in handle(self, op, args, kwargs)
105         isinstance(x, keras_tensor.KerasTensor)
106         for x in tf.nest.flatten([args, kwargs])):
---> 107       return TFOpLambda(op)(*args, **kwargs)
108     else:
109       return self.NOT_SUPPORTED

/usr/local/lib/python3.7/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
65     except Exception as e:  # pylint: disable=broad-except
66       filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67       raise e.with_traceback(filtered_tb) from None
68     finally:
69       del filtered_tb

TypeError: Exception encountered when calling layer "tf.math.truediv_3" (type TFOpLambda).

`x` and `y` must have the same dtype, got tf.float16 != tf.float32.

Call arguments received by layer "tf.math.truediv_3" (type TFOpLambda):   • x=tf.Tensor(shape=(None, None, None, 3), dtype=float16)   • y=tf.Tensor(shape=(3,), dtype=float32)   • name=None

解决方案 #

方案1：手动替换预处理逻辑，确保dtype匹配 #

关闭模型内置的预处理，自行实现预处理步骤并统一dtype：

from tensorflow.keras import layers
import tensorflow as tf
from tensorflow.keras.mixed_precision import set_global_policy
from tensorflow.keras.applications import efficientnet

set_global_policy(policy="mixed_float16")
input_shape = (224, 224, 3)
class_names = ["your_class_1", "your_class_2"]  # 替换为实际类别

# 关闭内置预处理，后续手动实现
base_model = efficientnet.EfficientNetB0(include_top=False, preprocess_input=False)
base_model.trainable = False

# 自定义预处理层，统一常量与输入的dtype
def custom_preprocess(input_tensor):
    # 执行EfficientNet标准预处理
    x = efficientnet.preprocess_input(input_tensor)
    # 将标准差常量转换为float16，匹配输入张量dtype
    stddev = tf.math.sqrt(tf.constant(efficientnet.IMAGENET_STDDEV_RGB, dtype=tf.float16))
    x = x / stddev
    return x

# 构建完整模型
inputs = layers.Input(shape=input_shape, name="input_layer")
x = custom_preprocess(inputs)
x = base_model(x, training=False)
x = layers.GlobalAveragePooling2D(name="pooling_layer")(x)
x = layers.Dense(len(class_names))(x)
# 输出层保持float32避免精度损失
outputs = layers.Activation("softmax", dtype=tf.float32, name="softmax_float32")(x)
model = tf.keras.Model(inputs, outputs)

model.compile(loss="sparse_categorical_crossentropy",
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

方案2：修改内置常量的dtype，适配混合精度策略 #

直接修改EfficientNet内置的标准差常量 dtype，使其与混合精度策略匹配：

from tensorflow.keras import layers
import tensorflow as tf
from tensorflow.keras.mixed_precision import set_global_policy
from tensorflow.keras.applications import efficientnet

set_global_policy(policy="mixed_float16")
input_shape = (224, 224, 3)
class_names = ["your_class_1", "your_class_2"]  # 替换为实际类别

# 将内置标准差常量转换为float16
efficientnet.IMAGENET_STDDEV_RGB = tf.cast(efficientnet.IMAGENET_STDDEV_RGB, dtype=tf.float16)

# 正常实例化模型即可
base_model = efficientnet.EfficientNetB0(include_top=False)
base_model.trainable = False

# 构建模型
inputs = layers.Input(shape=input_shape, name="input_layer")
x = base_model(inputs, training=False)
x = layers.GlobalAveragePooling2D(name="pooling_layer")(x)
x = layers.Dense(len(class_names))(x)
outputs = layers.Activation("softmax", dtype=tf.float32, name="softmax_float32")(x)
model = tf.keras.Model(inputs, outputs)

model.compile(loss="sparse_categorical_crossentropy",
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

方案3：通过Monkey Patch自动适配策略 #

通过替换EfficientNet的构建函数，自动将常量转换为当前混合精度策略的 dtype：

from tensorflow.keras import layers
import tensorflow as tf
from tensorflow.keras.mixed_precision import set_global_policy
from tensorflow.keras.applications import efficientnet

set_global_policy(policy="mixed_float16")
input_shape = (224, 224, 3)
class_names = ["your_class_1", "your_class_2"]  # 替换为实际类别

# 保存原始构建函数
original_efficientnet = efficientnet.EfficientNet

def patched_efficientnet(*args, **kwargs):
    # 获取当前全局策略的计算dtype
    policy = tf.keras.mixed_precision.global_policy()
    # 转换标准差常量为策略指定的dtype
    efficientnet.IMAGENET_STDDEV_RGB = tf.constant(efficientnet.IMAGENET_STDDEV_RGB, dtype=policy.compute_dtype)
    return original_efficientnet(*args, **kwargs)

# 替换原构建函数
efficientnet.EfficientNet = patched_efficientnet
# 适配EfficientNetB0的快捷调用
efficientnet.EfficientNetB0 = lambda **kwargs: patched_efficientnet(
    width_coefficient=1.0, depth_coefficient=1.0, default_size=224, **kwargs
)

# 正常实例化模型
base_model = efficientnet.EfficientNetB0(include_top=False)
base_model.trainable = False

# 构建模型
inputs = layers.Input(shape=input_shape, name="input_layer")
x = base_model(inputs, training=False)
x = layers.GlobalAveragePooling2D(name="pooling_layer")(x)
x = layers.Dense(len(class_names))(x)
outputs = layers.Activation("softmax", dtype=tf.float32, name="softmax_float32")(x)
model = tf.keras.Model(inputs, outputs)

model.compile(loss="sparse_categorical_crossentropy",
              optimizer=tf.keras.optimizers.Adam(),
              metrics=["accuracy"])

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