Error when trying to fit model - Tensorflow CNN

Question

I'm trying to create a CNN for image classification (Cats & Dogs). Everything was working fine untill the moment I used fit. Since I am a beginner I fear I haven't correctly created my sequential model. I'm also not sure if my steps per epoch and validation steps are correct.
I get the following error when I use history = model.fit:

2022-07-12 12:10:04.374122: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2)
Epoch 1/15
2022-07-12 11:57:57.465865: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2)>
2022-07-12 11:57:58.808025: I tensorflow/stream_executor/cuda/cuda_dnn.cc:369] Loaded cuDNN version 8201
2022-07-12 11:58:00.536214: W tensorflow/core/framework/op_kernel.cc:1680] Invalid argument: required broadcastable shapes
2022-07-12 11:58:00.536457: W tensorflow/core/framework/op_kernel.cc:1680] Invalid argument: required broadcastable shapes
2022-07-12 11:58:00.536536: W tensorflow/core/framework/op_kernel.cc:1680] Invalid argument: required broadcastable shapes
Traceback (most recent call last):
  File "C:\Users\myPC\PycharmProjects\pythonProject\catsdogs.py", line 75, in <module>
    history = model.fit(train_data_gen, epochs=epochs, validation_data=val_data_gen, steps_per_epoch=int(np.ceil(train_data_gen.n / float(batch_size))), validation_steps=int(np.ceil(val_data_gen.n / float(batch_size))))
  File "C:\Users\myPC\anaconda3\envs\tensorflow\lib\site-packages\keras\engine\training.py", line 1184, in fit
    tmp_logs = self.train_function(iterator)
  File "C:\Users\myPC\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\def_function.py", line 885, in __call__
    result = self._call(*args, **kwds)
  File "C:\Users\myPC\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\def_function.py", line 950, in _call
    return self._stateless_fn(*args, **kwds)
  File "C:\Users\myPC\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\function.py", line 3039, in __call__
    return graph_function._call_flat(
  File "C:\Users\myPC\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\function.py", line 1963, in _call_flat
    return self._build_call_outputs(self._inference_function.call(
  File "C:\Users\myPC\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\function.py", line 591, in call
    outputs = execute.execute(
  File "C:\Users\myPC\anaconda3\envs\tensorflow\lib\site-packages\tensorflow\python\eager\execute.py", line 59, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError:  required broadcastable shapes
     [[node Equal (defined at \PycharmProjects\pythonProject\catsdogs.py:75) ]] [Op:__inference_train_function_733]

Function call stack:
train_function

2022-07-12 11:58:00.649539: W tensorflow/core/kernels/data/generator_dataset_op.cc:107] Error occurred when finalizing GeneratorDataset iterator: Failed precondition: Python interpreter state is not initialized. The process may be terminated.
     [[{{node PyFunc}}]]

Process finished with exit code 1>

This is the code I use(Tensorflow version is 2.6):

import tensorflow as tf
import keras

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPool2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator

import os
import numpy as np
import matplotlib.pyplot as plt

# GET PROJECT FILES
PATH = 'cats_and_dogs'

train_dir = os.path.join(PATH, 'train')
validation_dir = os.path.join(PATH, 'validation')
test_dir = os.path.join(PATH, 'test')

# GET NUMBER OF FILES IN EACH DIRECTORY
total_train = sum([len(files) for r, d, files in os.walk(train_dir)])
total_val = sum([len(files) for r, d, files in os.walk(validation_dir)])
total_test = len(os.listdir(test_dir))

# VARIABLES FOR PRE-PROCESSING AND TRAINING.
batch_size = 128
epochs = 15
IMG_HEIGHT = 150
IMG_WIDTH = 150

# CREATE IMAGE DATA GENERATORS
train_image_generator = ImageDataGenerator(rotation_range=0.5, zoom_range=0.2, horizontal_flip=True, vertical_flip=True,
                                           rescale=1. / 255)
validation_image_generator = ImageDataGenerator(rescale=1. / 255)
test_image_generator = ImageDataGenerator(rescale=1. / 255)

train_data_gen = train_image_generator.flow_from_directory(directory=train_dir, target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                           class_mode='binary', batch_size=batch_size)
val_data_gen = validation_image_generator.flow_from_directory(directory=validation_dir,
                                                              target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                              class_mode='binary', batch_size=batch_size)
test_data_gen = test_image_generator.flow_from_directory(directory=test_dir, target_size=(IMG_HEIGHT, IMG_WIDTH),
                                                         class_mode='binary', batch_size=batch_size,
                                                         shuffle=False)


# CREATE MODEL
model = Sequential(
    [
        Conv2D(32, (3, 3), input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)),
        MaxPool2D((2, 2)),
        Dense(1, activation='relu')
    ]
)

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
              loss=tf.keras.losses.BinaryCrossentropy(),
              metrics=['accuracy'])

model.summary()
history = model.fit(train_data_gen, epochs=epochs, validation_data=val_data_gen, steps_per_epoch=int(np.ceil(train_data_gen.n / float(batch_size))), validation_steps=int(np.ceil(val_data_gen.n / float(batch_size))))

I feel like required broadcastable shapes is the culprit but I might be wrong.

Answer 1

After a bit of research, I finally made it work in pyCharm as well. The problem was I hadn't used Flatten. Here is the piece of my code that I changed. It works fine now. I'm still not sure why it wouldn't give me an error in Colab though.

model = Sequential(
    [
        keras.Input(shape=(IMG_WIDTH, IMG_HEIGHT, 3)),
        Conv2D(32, 3, padding='same', activation='relu'),
        MaxPooling2D(),
        Conv2D(64, 3, activation='relu'),
        MaxPooling2D(),
        Conv2D(128, 3, activation='relu'),
        Flatten(),
        Dense(64, activation='relu'),
        Dropout(0.5),
        Dense(1, activation='sigmoid')
    ]
)

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=3e-4),
    loss=keras.losses.BinaryCrossentropy(),
    metrics=["accuracy"]
)

model.summary()
history = model.fit(train_data_gen, epochs=epochs, validation_data=val_data_gen, batch_size=batch_size, verbose=2)