我是Tensorflow的新手，我知道有很多关于形状问题的帖子，但我还没有能够将他们的解决方案应用到我的问题上。所以，请原谅我，如果这是一个共同的/多余的问题。
上下文：我正在使用图像数据来预测图像的类型。更具体地说，我有许多图像的像素数据（50 x50 x3），这些图像分为4个类别[风景，肖像，抽象，其他]。
代码详细信息：我将尝试显示可能有问题的代码。当然，在这个项目中还有其他代码，但为了简单起见，我将省略它（否则它将是一个难以阅读的帖子）;然而，在后续工作中，如果没有找到解决方案，我将更详细地介绍。
我使用以下代码对标签进行了one-hot编码
train_labels = to_categorical(train_labels, num_classes=NUM_CLASSES, dtype='float32')
然后，我通过提供一个Tensor对象和一个图像路径列表来加载像素数据，并在所有路径上运行load_image函数（未显示）。

train_data = tf.data.Dataset.from_tensor_slices((train_files, train_labels)) 
train_data = train_data.map(load_image, num_parallel_calls=AUTOTUNE) 
train_data = train_data.map(normalize, num_parallel_calls=AUTOTUNE) 
train_data = train_data.shuffle(buffer_size=shuffle_buffer_size) 
train_data = train_data.batch(batch_size)`

train_data对象如下所示：
<_BatchDataset element_spec=(TensorSpec(shape=(None, 50, 50, 3), dtype=tf.float32, name=None), TensorSpec(shape=(None, 4), dtype=tf.float32, name=None))>
我相信这是有意义的，因为输入数据是50 x50 x3（具有批处理维度），输出维度是4，因为4个类的独热编码（具有批处理维度）。整个过程和结果与我的验证数据相同。
然而，当我尝试训练我的模型时，有一个形状错误：

def simple_FFNN(image_height, image_width, num_channels, num_classes):
    input_shape = \[image_height, image_width, num_channels\] # height, width, channels
    model = Sequential()
    model.add(layers.Input(shape=input_shape))
    model.add(layers.Dense(64, activation='relu'))
    model.add(layers.Dense(units=num_classes, activation='softmax'))

    return model

learning_rate = 0.01
epochs = 3
model = simple_FFNN(IMAGE_HEIGHT, IMAGE_WIDTH, NUM_CHANNELS, NUM_CLASSES)
optimizer = optimizers.SGD(learning_rate=learning_rate)
loss = losses.categorical_crossentropy
model.compile(loss=loss,
optimizer=optimizer,
metrics='accuracy')

training_results = model.fit(
    train_data,
    validation_data=validation_data,
    epochs=epochs,
    verbose=1)`

错误是：
ValueError: Shapes (None, 4) and (None, 50, 50, 4) are incompatible
第一个形状是我的输出形状，第二个形状几乎是我的输入形状，但通道数量略有减少。几乎就像班级的数量被替换了。
无论如何，我不明白为什么输出形状和输入形状必须首先匹配。我一定是误解了这个问题。有人有什么主意吗？
谢谢
我试着删除train_data.batch()中的字符串，但是我得到了一个错误，说（None，50，50，3）和（50，50，3）不匹配（这对我来说是有意义的）。有趣的是，它显示的是（50，50，3）而不是（50，50，4），所以我不明白为什么最后一个维度被换成了4。