为何keras.Model.fit()会将布尔张量转换为float32张量？

问题描述 #

构建自定义Keras模型时，输入为包含float32特征张量和布尔掩码张量的元组。直接调用模型时布尔张量类型正常，但将输入放入Dataset并使用fit()方法训练后，布尔张量被转换为float32类型，不符合掩码使用需求。

环境版本 #

• TensorFlow 2.15.0
• Keras 2.15.0
• Python 3.11.5

可复现代码 #

import tensorflow as tf
from tensorflow import keras

# 定义子类化模型
class TM(keras.Model):
    def __init__(self):
        super(TM, self).__init__()

    def call(self, inputs):
        features = inputs[0]
        input_mask = inputs[1]  # 此处通过fit调用时已被转为float32，预期应为bool
        print("features:")
        print(features)
        print()

        print("input mask:")
        print(input_mask)
        print("-"*100)
        return features

# 实例化模型并编译   
tm = TM()
tm.compile(loss="mean_squared_error")

# 创建测试数据
feature_tensor = tf.random.uniform(shape=[1,8])
boolean_tensor = tf.fill(dims=[1, 8], value=True)

# 构建TensorFlow Dataset
tds = tf.data.Dataset.from_tensor_slices((feature_tensor, boolean_tensor))

# 查看Dataset元素类型（此时类型正确）
print("ds element spec:")
print(tds.element_spec)
print("-"*100)

# 直接调用模型（类型正常）
print("直接调用模型的输入情况:")
print(tm((feature_tensor, boolean_tensor)))
print("-"*100)

# 通过Dataset和fit调用模型（布尔张量被转成float32）
tm.fit(tds)

输出结果 #

ds element spec:
(TensorSpec(shape=(8,), dtype=tf.float32, name=None), TensorSpec(shape=(8,), dtype=tf.bool, name=None))
----------------------------------------------------------------------------------------------------
直接调用模型的输入情况:
features:
tf.Tensor(
[[0.7957325  0.694777   0.1260978  0.5917362  0.17017555 0.8500334
  0.23767495 0.88851535]], shape=(1, 8), dtype=float32)

input mask:
tf.Tensor([[ True  True  True  True  True  True  True  True]], shape=(1, 8), dtype=bool)
----------------------------------------------------------------------------------------------------
tf.Tensor(
[[0.7957325  0.694777   0.1260978  0.5917362  0.17017555 0.8500334
  0.23767495 0.88851535]], shape=(1, 8), dtype=float32)
----------------------------------------------------------------------------------------------------
features:
Tensor("tm_7/strided_slice:0", shape=(), dtype=float32)

input mask:
Tensor("tm_7/strided_slice_1:0", shape=(), dtype=float32)
----------------------------------------------------------------------------------------------------
1/1 [==============================] - 0s 66ms/step - loss: 0.0417

原因分析 #

1. Keras输入自动类型对齐：fit()流程中，Keras会自动对齐输入数据类型以匹配损失函数或模型输出的类型。这里使用的mean_squared_error损失是针对float类型设计的，因此Keras会将所有输入张量统一转为float32。
2. Dataset预处理管道差异：直接调用模型时，输入数据会直接传入call方法；而通过fit()调用时，Dataset会经过Keras内置的输入预处理管道，该管道默认会将非float类型的张量转换为float类型，以保证训练流程的类型一致性。

解决方案 #

方案1：在模型call方法中显式转换类型 #

在call方法里把被转换的mask张量重新转回bool类型：

def call(self, inputs):
    features = inputs[0]
    input_mask = tf.cast(inputs[1], tf.bool)  # 显式转回bool
    # 后续掩码逻辑
    return features

方案2：在Dataset中强制保留bool类型 #

通过map操作确保mask张量始终保持bool类型，覆盖Keras的自动转换：

tds = tf.data.Dataset.from_tensor_slices((feature_tensor, boolean_tensor))
tds = tds.map(lambda x, y: (x, tf.cast(y, tf.bool)))  # 强制保留bool类型

方案3：用Input层显式定义输入类型 #

通过Keras的Input层明确指定每个输入的dtype，让模型严格按照定义接收输入：

# 显式定义输入类型
input_features = keras.Input(shape=(8,), dtype=tf.float32)
input_mask = keras.Input(shape=(8,), dtype=tf.bool)

# 构建模型
x = input_features  # 此处可添加自定义逻辑
model = keras.Model(inputs=[input_features, input_mask], outputs=x)

model.compile(loss="mean_squared_error")
model.fit(tds)

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