TensorFlow model.fit无法使用自定义Dataset，抛出未知秩形状错误

错误信息 #

<ipython-input-20-01755a6ded38> in <cell line: 1>()
----> 1 model.fit(
      2     dataset,
      3     epochs=100,
      4     verbose=1,
      5     batch_size=8)

1 frames /usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py in error_handler(*args, **kwargs)
    120             # To get the full stack trace, call:
    121             # `keras.config.disable_traceback_filtering()`
---> 122             raise e.with_traceback(filtered_tb) from None
    123         finally:
    124             del filtered_tb

/usr/local/lib/python3.10/dist-packages/keras/src/losses/loss.py in squeeze_or_expand_to_same_rank(x1, x2, expand_rank_1=True)
    105 def squeeze_or_expand_to_same_rank(x1, x2, expand_rank_1=True):
    106     """Squeeze/expand last dim if ranks differ from expected by exactly 1."""
---> 107     x1_rank = len(x1.shape)
    108     x2_rank = len(x2.shape)
    109     if x1_rank == x2_rank:

ValueError: Cannot take the length of shape with unknown rank.

复现代码 #

模型定义 #

model = keras.models.Sequential(
    [keras.layers.Dense(10, activation='relu'),
     keras.layers.Dense(1, activation='sigmoid')]
)
model.compile(loss='binary_crossentropy', metrics=['accuracy'])

数据集生成代码 #

# Random dataset
a = tf.convert_to_tensor(np.random.randint(0,100, size=[10,10]))
    
# Data generator class
class DataGenerator:
   def __init__(self, data, ratio=3):
      self._ratio = ratio
      self._data = data
    
   def __call__(self):
      shape = tf.shape(self._data).numpy()
      x = tf.convert_to_tensor(np.random.randint(1000,100000, size=[shape[0] * self._ratio, shape[1]]))
      x = tf.concat([self._data, x], axis = 0)
      y = tf.convert_to_tensor(np.random.random(shape[0]*(1 + self._ratio)))
    
      yield x, y


data_gen = DataGenerator(a, 3)
dataset = tf.data.Dataset.from_generator(
            data_gen,
            output_signature=(
                tf.TensorSpec(shape=(None,10), dtype=tf.int32),
                tf.TensorSpec(shape=(None), dtype=tf.float32)))

触发错误的训练代码 #

model.fit(
    dataset,
    epochs=100,
    verbose=1,
    batch_size=8)

解决方案 #

问题根源有两个：

1. 输出签名中y的形状定义错误：tf.TensorSpec(shape=(None), dtype=tf.float32) 会创建一个未知秩的张量，而Keras损失函数需要明确知道张量的秩才能正常计算。正确的一维张量形状应该写成(None,)，明确表示这是秩为1的张量。
2. tf.data.Dataset的批次处理方式错误：当向model.fit传入tf.data.Dataset时，batch_size参数会被忽略，必须直接在数据集上调用.batch()方法来实现分批次。

修改后的代码 #

修正数据集输出签名并添加批次处理

data_gen = DataGenerator(a, 3)
dataset = tf.data.Dataset.from_generator(
            data_gen,
            output_signature=(
                tf.TensorSpec(shape=(None,10), dtype=tf.int32),
                # 修正y的形状为(None,)，明确秩为1
                tf.TensorSpec(shape=(None,), dtype=tf.float32)))

# 在数据集上添加批次处理
dataset = dataset.batch(8)

修正后的训练代码

model.fit(
    dataset,
    epochs=100,
    verbose=1)

补充说明 #

直接传入生成器的x和y能正常训练，是因为Keras可以直接从生成的张量中推断出明确的秩；而通过tf.data.Dataset时，必须在输出签名中明确指定张量的秩，否则Keras无法正确解析张量形状，进而触发损失函数中的秩检测错误。

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