我正在用TensorFlow进行混合精度训练，以在GPU上加速训练。我启用了全局混合精度策略并实现了一个自定义损失函数。然而，在训练开始时，我收到了一个数据类型不匹配的错误。

错误：

TypeError: Input 'y' of 'Sub' Op has type float16 that does not match type float32 of argument 'x'

代码：

import tensorflow as tf
from tensorflow.keras import mixed_precision

# Enable mixed precision
mixed_precision.set_global_policy("mixed_float16")

# Custom loss function
def custom_mse(y_true, y_pred):
    diff = y_true - y_pred
    return tf.reduce_mean(tf.square(diff))

# Simple model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation="relu"),
    tf.keras.layers.Dense(1)
])

model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss=custom_mse,
)

# Dummy data
import numpy as np
x = np.random.rand(100, 10)
y = np.random.rand(100, 1)

model.fit(x, y, epochs=3)

模型能够正确编译，但在训练阶段因上述数据类型不匹配的错误而失败。

我的问题：

1. 为什么在启用混合精度时会出现数据类型不匹配？
2. 如何在混合精度下实现自定义损失函数？

解决方案

When using the mixed_float16 policy, your model outputs predictions (y_pred) as float16, but your true target data (y_true) stays as float32. Since TensorFlow's arithmetic operations need the data types to match exactly, you cannot perform operations between a float32 and a float16 operand.

You need to explicitly cast both inputs to float32 in your custom loss function before doing any calculations. Also, to ensure numerical stability and to avoid underflow, you must make your model's final output layer a float32。

import tensorflow as tf
from tensorflow.keras import mixed_precision
import numpy as np

mixed_precision.set_global_policy("mixed_float16")

# Custom loss function
def custom_mse(y_true, y_pred):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)

    diff = y_true - y_pred
    return tf.reduce_mean(tf.square(diff))

model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation="relu"),
    tf.keras.layers.Dense(1, dtype='float32') 
])

model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss=custom_mse,
)

x = np.random.rand(100, 10)
y = np.random.rand(100, 1)

model.fit(x, y, epochs=3)