如何修复使用tf.data.Dataset训练CNN时Dense层的ValueError？

使用自定义Model类构建CNN，训练数据采用tf.data.Dataset格式，训练时从卷积层切换到Dense层抛出ValueError，提示Dense层输入的最后维度未定义（发现None），输入形状为(None, None, None, 3)。使用CIFAR10数据集时正常，自定义数据集出现该问题。

代码概要及报错信息如下：

class CNN(Model):
    def __init__(self, input_shape):
        super(CNN, self).__init__()
        conv_arguments = dict(kernel_size = (3,3), activation = 'relu')
        pool_arguments = dict(pool_size = (2,2))

        self.conv1 = Conv2D(32, **conv_arguments, input_shape=input_shape)
        self.maxpool1 = MaxPooling2D(**pool_arguments)
        self.conv2 = Conv2D(48, **conv_arguments)
        self.maxpool2 = MaxPooling2D(**pool_arguments)
        self.conv3 = Conv2D(64, **conv_arguments)
        self.maxpool3 = MaxPooling2D(**pool_arguments)
        
        self.flatten = Flatten()
        self.dense1 = Dense(64, activation = 'relu')
        self.dropout = Dropout(0.2)
        self.dense2 = Dense(4, activation = 'sigmoid')

    def call(self, x):
        x = self.conv1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.maxpool2(x)
        x = self.conv3(x)
        x = self.maxpool3(x)
        x = self.flatten(x)
        x = self.dense1(x)
        x = self.dropout(x)
        x = self.dense2(x)
        return x

# 数据集构建
train_data = tf.data.Dataset.from_tensor_slices((train_files, train_labels))
train_data = train_data.shuffle(len(train_files))
train_data = train_data.map(lambda x, y: (load_image(x), y), num_parallel_calls=AUTOTUNE)
train_data = train_data.cache().repeat()
train_data = train_data.map(lambda x, y: (augmentation.call(x), y), num_parallel_calls=AUTOTUNE)
train_data = train_data.batch(batch_size)
train_data = train_data.prefetch(AUTOTUNE)

# 模型训练
model = CNN(input_shape=(424, 424, 3))
model.compile(optimizer=opt, loss=loss_obj, metrics=['accuracy'])
hist = model.fit(train_data, epochs=epochs, verbose = 1, callbacks = [early], validation_data = val_data, class_weight = data.class_weights, steps_per_epoch = nb_steps_per_epoch)

报错信息：

ValueError: The last dimension of the inputs to a Dense layer should be defined. Found None. Full input shape received: (None, None)

Call arguments received by layer "cnn" (type CNN):
  • x=tf.Tensor(shape=(None, None, None, 3), dtype=float32)

核心问题是输入数据的形状未固定：

• CIFAR10数据集的图片尺寸统一（32x32），模型能明确后续层的输出形状；
• 自定义数据集的原始图片尺寸不一致，且load_image或augmentation过程中未强制统一到固定尺寸(424,424,3)，导致模型接收的输入形状为(None, None, None, 3)（批量维度+不确定的高/宽/通道），卷积层处理后Flatten的输出最后维度为None，Dense层无法处理动态形状的输入。

1. 强制统一图片尺寸 #

修改load_image函数，确保所有图片加载后缩放到固定尺寸(424,424)：

def load_image(file_path):
    img = tf.io.read_file(file_path)
    img = tf.image.decode_image(img, channels=3)
    # 强制缩放到固定尺寸
    img = tf.image.resize(img, (424, 424))
    img = tf.cast(img, tf.float32) / 255.0  # 归一化（可选）
    return img

2. 验证增强操作的形状一致性 #

确保数据增强操作不会改变图片的固定尺寸，例如：

• 使用tf.image.random_crop时，指定size=(424, 424, 3)；
• 避免使用可能导致尺寸变化的操作，或在操作后重新resize到固定尺寸。

示例增强函数：

def augmentation(img):
    img = tf.image.random_flip_left_right(img)
    # 随机裁剪后保持尺寸
    img = tf.image.random_crop(img, size=(424, 424, 3))
    img = tf.image.random_brightness(img, max_delta=0.2)
    return img

3. 检查数据集的实际形状 #

在训练前验证数据集的输出形状，确保每个batch的图片形状为(batch_size, 424, 424, 3)：

for x_batch, y_batch in train_data.take(1):
    print("Batch shape:", x_batch.shape)
    # 输出应为 (batch_size, 424, 424, 3)

如果输出不符合预期，说明load_image或augmentation存在问题，需要进一步排查。

4. 显式构建模型（辅助步骤） #

在模型编译前显式调用build方法，固定输入形状：

model = CNN(input_shape=(424, 424, 3))
model.build(input_shape=(None, 424, 424, 3))  # None表示批量维度
model.summary()  # 查看各层输出形状是否正常
model.compile(...)

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