使用tf.datasets的input_fn在GCMLE上1轮训练后报错求助

你的猜测完全正确！这个报错确实是因为训练到第一轮结束时，剩余的样本数量不足设定的batch_size（110条），最后一批只拿到了74条数据，而你的损失计算逻辑（从报错看是稀疏softmax交叉熵）没有处理这种可变batch大小的情况，最终引发了形状不匹配的错误。

为什么旧的队列式input_fn没问题？ #

旧的基于队列的输入管道通常默认会丢弃最后一批不足batch_size的数据，或者在队列逻辑中自动处理了这种情况，所以不会触发形状不匹配的报错；而你当前使用的tf.contrib.data.map_and_batch默认是保留最后一批数据的（drop_remainder参数默认值为False），这就导致了最后一批的形状和之前的批次不一致。

解决方法 #

最直接且通用的解决方案是在map_and_batch中设置drop_remainder=True，让TensorFlow自动丢弃最后一批不足指定batch大小的数据：

修改你的input_fn中对应的代码行：

dataset = dataset.apply(tf.contrib.data.map_and_batch(
    lambda row: parse_csv_dataset(row, hparams = hparams), 
    batch_size = batch_size, 
    num_parallel_batches = multiprocessing.cpu_count(),
    drop_remainder=True  # 新增这个参数
))

这个参数会确保每一批数据的大小严格等于batch_size，彻底避免形状不匹配的问题，同时训练轮数的计算也会更准确（每一轮都是整数个完整批次）。

可选优化（非报错直接原因，但建议调整） #

另外注意你当前代码中shuffle、repeat和prefetch的顺序，正确的流水线顺序应该是：

if shuffle:
    dataset = dataset.shuffle(buffer_size = 10000)
dataset = dataset.repeat(num_epochs)
dataset = dataset.prefetch(1)

这样每一轮epoch的数据都会被独立打乱，而不是把多轮数据重复后再整体打乱，能保证训练数据的随机性，同时预取操作放在最后也能最大化流水线的效率。

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