Problem with input and output shapes in Flatten() layer in Tensorflow

Question

I am trying to build a basic Convolutional Neural Network for image classification. I have a dataset of images belonging to 4 classes. I created Tensorflow datasets using ImageDataGenerator and flow_from_directory :

train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale=1./255,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True,
    validation_split=0.2
)

val_datagen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1./255, validation_split=0.2)

train_data = train_datagen.flow_from_directory(
    dataset_source, 
    color_mode='rgb',
    batch_size=batch_size, 
    class_mode='categorical',
    shuffle=True,
    subset='training'
) 

val_data = val_datagen.flow_from_directory(
    dataset_source, 
    color_mode='rgb',
    batch_size=batch_size, 
    class_mode='categorical',
    shuffle=False,
    subset='validation'
)

An example of architecture of my CNN that I tested is trivial and looks as follows:

image_shape = (299, 299, 3)

model = tf.keras.Sequential()
model.add(tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=image_shape))
model.add(tf.keras.layers.MaxPooling2D((2, 2)))
model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))
model.add(tf.keras.layers.MaxPooling2D((2, 2)))
model.add(tf.keras.layers.Conv2D(128, (3, 3), activation='relu'))

model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(128, activation='relu'))
model.add(tf.keras.layers.Dense(4))

model.compile(loss = "categorical_crossentropy",
             optimizer = "adam",
             metrics = ["accuracy"])

Whenever I try to start training using model.fit(train_data, epochs=100, validation_data = val_data) , I receive the following error related to Flatten() layer:

InvalidArgumentError:  Input to reshape is a tensor with 7372800 values, but the requested shape requires a multiple of 645248
[[node sequential_1/flatten_1/Reshape (defined at <ipython-input-9-d01204576b1d>:1) ]] [Op:__inference_train_function_1823]

I checked it using various versions of CNN, even the simplest containing single convolution layers, but the error occurs every time, only with different values. It's probably quite obvious but despite many tries I am not able to solve it, so I will be very thankful for any directions how to deal with my issue.

Answer 1

Using flow_from_directory calls without specifying a target_size is just asking for trouble, because even a single wrongly-sized input will mess up your model fitting process. Simple fix is to add target_size=(299, 299) to the calls just to be safe (as seen in comment).

Also, consider adding image_shape = (299, 299, 3) above the flow_from_directory calls and specifying target_size = image_shape[:-1] to just pull it from there. This will make modifying image_shape easier, which you may want to do to experiment with various input size scalings.