TensorFlow单独训练顶层时logits与labels尺寸不匹配报错如何解决

你遇到的维度不匹配错误核心是输入数据与模型预期形状不符：
你当前定义的model仅为顶层分类头，预期输入是基础模型输出的(7,7,512)维度瓶颈特征，但你用flow_from_directory生成的是(224,224,3)维度的原始图片数据，直接喂给分类头会导致前向传播时维度错位，最终出现logits和标签维度不匹配的报错。
batch_size设为1时能跑通属于维度巧合，但输入数据完全不符合分类头的预期，因此训练得到的模型准确率极低。

根据你的训练需求，可任选以下两种方案实现大数据集的分步训练：

• 方案1：拼接基础模型与分类头，直接训练图片数据
  将你之前用来生成瓶颈特征的预训练基础模型（如VGG16、ResNet的卷积层部分）和现有分类头拼接为完整模型，冻结基础模型权重后直接用图片生成器训练，无需提前生成特征：

  from tensorflow.keras.applications import VGG16
  # 加载预训练基础模型，去掉自带的顶层分类头
  base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224,224,3))
  # 冻结基础模型权重，训练时不更新
  base_model.trainable = False
  
  # 拼接完整训练模型
  full_model = Sequential([
      base_model,
      # 原有分类头结构保持不变
      Flatten(input_shape=(7, 7, 512)),
      Dense(512, activation="relu"),
      Dropout(0.7),
      Dense(num_classes, activation='softmax')
  ])
  
  # 编译后直接用你定义的图片生成器训练即可
  full_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
  full_model.fit(
      train_generator,
      epochs=10,
      steps_per_epoch=train_generator.samples//train_generator.batch_size
  )
  

• 方案2：封装生成器实时生成瓶颈特征
  如果你不想修改现有分类头结构，可以对图片生成器做二次封装，每次生成batch时先调用基础模型生成瓶颈特征，再喂给分类头训练：

  def bottleneck_generator(original_generator, base_model):
      while True:
          imgs, labels = next(original_generator)
          # 实时生成瓶颈特征
          features = base_model.predict(imgs, verbose=0)
          yield features, labels
  
  # 初始化瓶颈特征生成器
  train_bottleneck_gen = bottleneck_generator(train_generator, base_model)
  # 用原有分类头训练
  model.fit(
      train_bottleneck_gen,
      epochs=10,
      steps_per_epoch=train_generator.samples//train_generator.batch_size
  )
  

注意：两种方案调用fit时都必须指定steps_per_epoch参数，值为训练集总样本数除以batch_size取整，避免生成器循环时出现维度异常。

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