TensorFlow Conv2D层报输入维度不兼容（期望4维实得3维）

报错的核心原因是tf.data.Dataset的所有变换操作（shuffle、repeat、batch等）都不会原地修改原数据集对象，执行后会返回一个全新的、应用了对应变换的Dataset实例。
你的input_fn中存在一处写法错误：

dataset.shuffle(SHUFFLE_SIZE).repeat(epochs).batch(batch_size)

这行代码执行了数据集打乱、重复、分批次的链式操作，但没有将返回的新数据集赋值回dataset变量。最终函数返回的还是最开始通过from_tensor_slices创建的、逐样本输出的原始数据集，Estimator拉取数据时拿到的单条样本形状为(28, 28, 1)，是3维结构，缺少Conv2D层要求的batch维度，因此触发维度不匹配报错。
你之前做的数据集reshape、模型输入层shape配置都是正确的，问题和模型结构、数据预处理无关。

只需要将数据集链式变换的返回值重新赋值给dataset变量即可，修正后的input_fn代码如下：

def input_fn(images, labels, epochs, batch_size):
    dataset = tf.data.Dataset.from_tensor_slices((images, labels))

    SHUFFLE_SIZE = 5000
    # 接收变换操作返回的新数据集
    dataset = dataset.shuffle(SHUFFLE_SIZE).repeat(epochs).batch(batch_size)
    dataset = dataset.prefetch(None)

    return dataset

修复后数据集输出的单批次数据形状为(BATCH_SIZE, 28, 28, 1)，完全匹配Conv2D层的输入维度要求，报错会直接消失。
可选优化：可以将prefetch的参数从None替换为tf.data.AUTOTUNE，让TensorFlow自动调整预取缓冲区大小，提升数据加载效率。

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