Keras/TensorFlow下CNN蘑菇分类任务model.fit()维度不匹配报错求助

• 核心为卷积池化输出的4维特征图未做展平，直接接入分类全连接层，导致模型输出维度与标签维度不匹配：
  • 数据集每个批次的标签形状为(32,)，对应批次大小32，每个样本对应1个分类标签
  • 当前模型最终输出形状为(None, 20, 20, 10)，TensorFlow计算损失时会自动将前三个维度拍平为(32*20*20=12800, 10)，和标签第一维度32不匹配，触发报错。

在ReLU激活层和最终全连接层之间加入Flatten()层，将202032的特征图展平为一维向量，再接入10分类的全连接层即可。也可使用GlobalAveragePooling2D()代替Flatten层，减少参数量，降低过拟合风险。

修改后的模型构建代码如下：

#Constructing layers
input_layer = keras.Input(shape=(64, 64, 3))
x = layers.Conv2D(filters=32, kernel_size=(3, 3), activation="relu")(input_layer)
x = layers.MaxPooling2D(pool_size=(3, 3))(x)
x = keras.layers.ReLU()(x)
# 新增展平层，将4维特征转换为匹配全连接层输入的2维结构
x = layers.Flatten()(x)
output = layers.Dense(10, activation="softmax")(x)

#Making and fitting the model
model = keras.Model(inputs=input_layer, outputs=output)
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=['accuracy'])
model.fit(train_set, epochs=5, validation_data=validation_set)

修改后模型输出形状变为(None, 10)，和标签的批次维度完全匹配，损失计算即可正常运行。

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