使用VGG16与tf.data.Dataset训练模型时输入形状不兼容问题

我尝试用TensorFlow结合预训练VGG16构建模型，但一直遇到输入形状不兼容的问题。和公开代码的唯一区别是我用tf.data.Dataset加载数据，而非tf.image的DirectoryIterator。

以下是我的代码：

zip_ref = ZipFile(zip_file, 'r')
zip_ref.extractall(repository_dir)
zip_ref.close()

train_dir = os.path.join(repository_dir, "seg_train", "seg_train")
test_dir = os.path.join(repository_dir, "seg_test", "seg_test")

os.system(f"rm -r {os.path.join(repository_dir, 'seg_pred')}")

# load variables
validation_percentage = 0.2
label_mode = "int"
# for our model purposes
img_size = (224, 224)
color_mode='rgb'

data_train, data_val = image_dataset_from_directory(
    train_dir,
    batch_size=None,
    label_mode=label_mode,
    color_mode=color_mode,
    image_size=img_size,
    validation_split=validation_percentage,
    subset="both",
    seed=123,
)
data_test = image_dataset_from_directory(
    test_dir,
    batch_size=None,
    label_mode=label_mode,
    color_mode=color_mode,
    image_size=img_size,
)

classes = data_train.class_names
print(classes)

scale = 1.0/255
normalization_layer = tf.keras.layers.Rescaling(scale)
data_train_norm = data_train.map(lambda x,y: (normalization_layer(x), y))
data_val_norm = data_val.map(lambda x,y: (normalization_layer(x), y))
data_test_norm = data_test.map(lambda x,y: (normalization_layer(x), y))

input_size = None
for img, label in data_train_norm.take(1).as_numpy_iterator():
    input_size = img.shape
print(input_size)

base_model = VGG16(
    input_shape=input_size, # Shape of our images
    include_top = False, # Leave out the last fully connected layer
    weights = 'imagenet'
)

# we do not train the parameters
for layer in base_model.layers:
    layer.trainable = False

# Flatten the output layer to 1 dimension
x = layers.Flatten()(base_model.output)

# FC layer very simple and with a softmax activation unit
x = layers.Dense(len(classes), activation="softmax")(x)

landscapeModel01 = Model(inputs=base_model.input, outputs=x, name="landscapeModel01")

loss = "sparse_categorical_crossentropy"
optimizer = "adam"

landscapeModel01.compile(
    optimizer=optimizer, 
    loss=loss,
    metrics=["loss","accuracy"]
)

#fit data
shuffle=True # variable
epochs=50 # variable, according if it is able to converge
batch_size = 200

print(landscapeModel01.input)

landscapeModel01.fit(
    data_train_norm,
    validation_data=data_val_norm,
    epochs=epochs,
    shuffle=shuffle,
    batch_size=batch_size
)

运行后出现如下错误：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In [10], line 8
      4 batch_size = 200
      6 print(landscapeModel01.input)
----> 8 landscapeModel01.fit(
      9     data_train_norm,
     10     validation_data=data_val_norm,
     11     epochs=epochs,
     12     shuffle=shuffle,
     13     batch_size=batch_size
     14 )

File ~/anaconda3/envs/faa/lib/python3.10/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /tmp/__autograph_generated_file8y_bf523.py:15, in outer_factory.<locals>.inner_factory.<locals>.tf__train_function(iterator)
     13 try:
     14     do_return = True
---> 15     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16 except:
     17     do_return = False

ValueError: in user code:

    File "/home/renan/anaconda3/envs/faa/lib/python3.10/site-packages/keras/engine/training.py", line 1160, in train_function  *
        return step_function(self, iterator)
    File "/home/renan/anaconda3/envs/faa/lib/python3.10/site-packages/keras/engine/training.py", line 1146, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/home/renan/anaconda3/envs/faa/lib/python3.10/site-packages/keras/engine/training.py", line 1135, in run_step  **
        outputs = model.train_step(data)
    File "/home/renan/anaconda3/envs/faa/lib/python3.10/site-packages/keras/engine/training.py", line 993, in train_step
        y_pred = self(x, training=True)
    File "/home/renan/anaconda3/envs/faa/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/home/renan/anaconda3/envs/faa/lib/python3.10/site-packages/keras/engine/input_spec.py", line 295, in assert_input_compatibility
        raise ValueError(

    ValueError: Input 0 of layer "landscapeModel01" is incompatible with the layer: expected shape=(None, 224, 224, 3), found shape=(224, 224, 3)

• 问题核心：加载数据时设置batch_size=None，导致数据集的每个元素是单张图片（形状为(224,224,3)），而Keras模型要求输入必须包含batch维度（预期形状为(None,224,224,3)，其中None代表batch大小）。
• 修改方案二选一即可：
  1. 加载数据集时直接指定batch_size（比如batch_size=32），替换代码中的batch_size=None
  2. 保留batch_size=None，但在数据归一化后为数据集添加batch维度：

     data_train_norm = data_train_norm.batch(32)
     data_val_norm = data_val_norm.batch(32)
     data_test_norm = data_test_norm.batch(32)
     

• 额外注意：使用tf.data.Dataset作为fit()输入时，无需再在fit()中指定batch_size参数，直接删除batch_size=batch_size即可。

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