TensorFlow使用from_generator构建数据集时形状不匹配报错问题

问题根因 #

两处配置不匹配导致报错和警告：

1. 初始output_shapes错误添加了batch维度：tf.data.Dataset.from_generator的output_shapes参数指定的是单条样本的形状，不需要提前包含batch维度（batch维度后续调用.batch()方法时会自动补充），你初始写的[None, 500]相当于要求单条样本就带batch维度，和生成器返回的单条样本形状(500,)冲突，触发第一个ValueError。
2. 生成器返回的输入张量多了多余维度：你修改output_shapes后触发的警告，说明data_generator.generate返回的ids1、ids2实际形状是(500, 1)而非预期的(500,)，多了最后一维的冗余维度，喂入模型时会被误识别为特征维度为1，和Input层定义的500维特征不匹配。

解决步骤 #

• 第一步：修正生成器返回的输入形状，去掉冗余的最后一维
  在预处理返回前添加维度压缩操作：

  # 去除最后一维的冗余维度，保证输入形状为(500,)
  ids1 = tf.squeeze(ids1, axis=-1)
  ids2 = tf.squeeze(ids2, axis=-1)
  return ({'input_1': ids1, 'input_2': ids2}, label)
  

• 第二步：修正数据集output_shapes配置，明确单条样本的形状

  ds = tf.data.Dataset.from_generator(
       data_generator.generate,
       output_types=({"input_1": tf.int64, "input_2": tf.int64}, tf.int64), 
       # 单条样本输入形状为(500,)，标签如果是标量则形状为[]
       output_shapes = ({"input_1": tf.TensorShape([500]), "input_2": tf.TensorShape([500])}, tf.TensorShape([]))
  )
  

• 第三步：按正常流程添加batch操作
  后续调用.batch(batch_size)方法时，tf.data会自动为所有张量补充batch维度，最终输入模型的形状为(batch_size, 500)，和Input层定义的(None, 500)完全匹配，不会再出现形状不匹配问题。

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