I am trying to make an image classifier. When running this it does okay on the first epoch, but at the beginning of the second epoch it runs out of training data, gives me a warning, then just starts the third epoch, and when thats done, it gives me another error on epoch 4, and so on.

I believe the problem is in the image data generator, because it doesnt restart after every epoch (i thought it always did that) but i didnt find any way to restart it after every epoch.

Here is the code:

import os
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, Dropout
from tensorflow.keras.optimizers import Adam
import numpy as np
import matplotlib.pyplot as plt
import shutil

os.environ["CUDA_VISIBLE_DEVICES"] = "-1" 
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

train = 'train'
toclassify = 'toclassify'
results = 'results'
resultsyes = os.path.join(results, 'yes')
resultsno = os.path.join(results, 'no')


os.makedirs(resultsyes, exist_ok=True)
os.makedirs(resultsno, exist_ok=True)


size = (128, 128)
batch_size = 32  
epochs = 10


datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    validation_split=0.2
)

train_generator = datagen.flow_from_directory(
    train,
    target_size=size,
    batch_size=batch_size,
    class_mode='binary',
    subset='training',
    color_mode='grayscale'
)

validation_generator = datagen.flow_from_directory(
    train,
    target_size=size,
    batch_size=batch_size,
    class_mode='binary',
    subset='validation',
    color_mode='grayscale'
)

print(f"for training: {train_generator.samples}")
print(f"for validation: {validation_generator.samples}")


steps = train_generator.samples // batch_size
steps_v = validation_generator.samples // batch_size #validation steps


model = Sequential([
    Conv2D(32,(3, 3),activation='relu',input_shape=(size[0],size[1],1)),
    MaxPooling2D((2,2)),
    Dropout(0.1),
    Conv2D(64,(3, 3),activation='relu'),
    MaxPooling2D((2,2)),
    Dropout(0.1),
    Conv2D(128,(3, 3),activation='relu'),
    MaxPooling2D((2,2)),
    Flatten(),
    Dense(512,activation='relu'),
    Dense(1,activation='sigmoid')
])

model.compile(optimizer='adam',loss='binary_crossentropy',metrics=['accuracy'])

history = model.fit(
    train_generator,
    steps_per_epoch=steps,
    validation_data=validation_generator,
    validation_steps=steps_v,
    epochs=epochs
)

Any help would be much much appreciated <3

I tried:

1. Changing the batch size
2. Adding/removing images
3. Subtracting 1 from steps and steps_v in hopes that it would restart before the next epoch
4. Looking for an answer to this question on google
5. Using ai tools

I expected:

1. I expected it to just dont give me the warning and skip whole epochs because it messes up the pyplot after that and i cant tell whether the model is overfitting or not and i dont think its normal for it to do that