I am facing a shape mismatch error during the training of a convolutional neural network (CNN) using TensorFlow and Keras.

I am using a custom LeNet-5 architecture defined as follows:

import tensorflow as tf
from keras import layers, models

def create_lenet5(input_shape, num_classes):
    model = models.Sequential([
        layers.Conv2D(6, kernel_size=(5, 5), strides=(1, 1), activation='relu', input_shape=input_shape),
        layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)),
        layers.Conv2D(16, kernel_size=(5, 5), strides=(1, 1), activation='relu'),
        layers.MaxPooling2D(pool_size=(2, 2), strides=(2, 2)),
        layers.Flatten(),
        layers.Dense(120, activation='relu'),
        layers.Dense(84, activation='relu'),
        layers.Dense(num_classes, activation='softmax')
    ])
    return model

The error message indicates that the labels have a shape of (240,), which implies that 240 labels are being passed instead of a smaller batch size.
The logits (predictions from the model) have a shape of (16, 14), indicating that the model is outputting predictions for a batch size of 16 and 14 classes.

data_preprocessing script:

import os
import cv2
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
import re

class DataPreprocessor:
    def __init__(self, image_dir, annotation_file, target_size=(256, 256)):
        self.image_dir = image_dir
        self.annotation_file = annotation_file
        self.target_size = target_size

        # Load annotations
        print("Loading annotations...")
        self.annotations = pd.read_csv(annotation_file, encoding='utf-8')
        print(f"Loaded {len(self.annotations)} annotations.")

    def load_and_preprocess_image(self, file_path):
        image = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
        if image is None:
            print(f"Failed to load image: {file_path}")
            return None
        image = cv2.resize(image, self.target_size)
        return image.astype(np.float32) / 255.0  # Normalize to [0, 1]

    def batch_generator(self, X, y, batch_size):
        num_samples = len(X)
        while True:
            indices = np.arange(num_samples)
            np.random.shuffle(indices)
            for start in range(0, num_samples, batch_size):
                batch_indices = indices[start:start + batch_size]
                batch_images = []
                batch_labels = []

                for idx in batch_indices:
                    image_path = X[idx]

                    if os.path.exists(image_path):
                        image = self.load_and_preprocess_image(image_path)
                        if image is not None:
                            batch_images.append(image)
                            label = self.process_labels(y[idx])
                            batch_labels.append(label)
                        else:
                            print(f"Processed image is None for: {image_path}")
                    else:
                        print(f"Image not found for: {image_path}. Skipping...")

                if batch_images and batch_labels:
                    batch_images = np.array(batch_images)  # Shape: (batch_size, 256, 256)
                    batch_labels = np.array(batch_labels)  # Shape: (batch_size, num_classes)

                    # Ensure correct shapes
                    if batch_images.ndim == 3:
                        batch_images = np.expand_dims(batch_images, axis=-1)  # Shape: (batch_size, 256, 256, 1)

                    # Ensure batch_labels are 2D
                    if batch_labels.ndim == 1:
                        batch_labels = np.expand_dims(batch_labels, axis=0)  # Shape: (1, num_classes)
                    
                    if batch_labels.ndim == 2 and batch_labels.shape[0] == 1:
                        batch_labels = np.repeat(batch_labels, len(batch_images), axis=0)

                    print(f"Yielding batch: images shape {batch_images.shape}, labels shape {batch_labels.shape}")  # Debug output

                    if batch_images.shape[0] == batch_labels.shape[0]:  # Ensure batch sizes match
                        yield batch_images, batch_labels
                    else:
                        print(f"Warning: Batch size mismatch - images: {batch_images.shape[0]}, labels: {batch_labels.shape[0]}")
                else:
                    print("Warning: Empty batch detected, skipping this batch...")

    def load_dataset_paths_and_labels(self):
        print("Loading dataset paths and labels...")
        image_paths = []
        labels = []
        total_images = len(self.annotations)

        for i in range(total_images):
            image_name = self.annotations.iloc[i]['Image Index']
            image_path = os.path.join(self.image_dir, image_name)
            if os.path.exists(image_path):
                image_paths.append(image_path)
                labels.append(self.annotations.iloc[i]['Finding Labels'])  # Store as is
            else:
                print(f"Image not found for {image_name}. Skipping...")

            # Print progress every 100 images
            if (i + 1) % 100 == 0 or (i + 1) == total_images:
                print(f"Loaded {i + 1}/{total_images} images and labels.")

        return image_paths, labels

    def process_labels(self, labels_str):
        if isinstance(labels_str, bytes):
            labels_str = labels_str.decode('utf-8')

        label_list = re.split(r'||,', labels_str)
        label_dict = {
            'Atelectasis': 0,
            'Cardiomegaly': 1,
            'Effusion': 2,
            'Infiltration': 3,
            'Mass': 4,
            'Nodule': 5,
            'Pneumonia': 6,
            'Pneumothorax': 7,
            'Consolidation': 8,
            'Edema': 9,
            'Emphysema': 10,
            'Fibrosis': 11,
            'Pleural_Thickening': 12,
            'Hernia': 13,
            'No Finding': 14
        }

        # Initialize an array for the one-hot encoded labels
        label_array = np.zeros((len(label_dict),), dtype=np.int32)

        # Process the labels
        for label in label_list:
            label = label.strip()
            if label in label_dict:
                label_array[label_dict[label]] = 1  # One-hot encoding

        return label_array

    def split_data(self, images, labels, test_size=0.3):
        print("Splitting data...")
        X_train, X_temp, y_train, y_temp = train_test_split(images, labels, test_size=test_size, random_state=42)
        X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=42)
        print(f"Data split complete: {len(X_train)} training samples, {len(X_val)} validation samples, {len(X_test)} test samples.")
        return X_train, X_val, X_test, y_train, y_val, y_test`

train:

import os
import tensorflow as tf
from utils.config import Config
from data_preprocessing import DataPreprocessor
from models.classification_model import create_lenet5

def main():
    # Initialize DataPreprocessor
    preprocessor = DataPreprocessor(Config.IMAGE_DIR, Config.ANNOTATION_FILE)

    if preprocessor.annotations is None or preprocessor.annotations.empty:
        print("Error: Annotations not loaded or empty. Exiting...")
        return

    # Load dataset paths and labels
    image_paths, labels = preprocessor.load_dataset_paths_and_labels()
    X_train, X_val, X_test, y_train, y_val, y_test = preprocessor.split_data(image_paths, labels, test_size=0.3)
    print(f"Data split complete: {len(X_train)} training samples, {len(X_val)} validation samples, {len(X_test)} test samples.")

    # Create LeNet-5 model
    model = create_lenet5(input_shape=preprocessor.target_size + (1,), num_classes=Config.NUM_CLASSES)

    # Compile the model
    model.compile(optimizer=tf.keras.optimizers.Adam(),
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                  metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

    # Define callbacks (optional)
    checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath='models/best_model.keras', save_best_only=True)
    early_stopping_callback = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)

    # Train the model using generators
    train_generator = preprocessor.batch_generator(X_train, y_train, batch_size=Config.BATCH_SIZE)

    # Training loop
    for epoch in range(Config.EPOCHS):
        print(f"Epoch {epoch + 1}/{Config.EPOCHS}")
        epoch_loss = 0
        epoch_accuracy = 0
        steps = 0

        # Debug: Check if the generator is yielding batches
        try:
            for batch_images, batch_labels in train_generator:
                if batch_images is None or batch_labels is None:
                    print("Received None for batch_images or batch_labels. Skipping...")
                    continue

                steps += 1
                print(f"Training on batch: images shape {batch_images.shape}, labels shape {batch_labels.shape}")

                # Train on the current batch
                history = model.train_on_batch(batch_images, batch_labels)

                epoch_loss += history[0]
                epoch_accuracy += history[1]

                print(f"Step {steps}: Loss = {history[0]}, Accuracy = {history[1]}")

                # Optionally, break after a few steps for testing
                if steps >= 10:  # Adjust this to control how many steps you want to test
                    break

        except Exception as e:
            print(f"Error during training: {e}")
            break

        avg_loss = epoch_loss / steps if steps > 0 else 0
        avg_accuracy = epoch_accuracy / steps if steps > 0 else 0
        print(f"Epoch {epoch + 1} completed: Avg Loss = {avg_loss}, Avg Accuracy = {avg_accuracy}")

    # Save the final model
    model.save('models/final_model.keras')

if __name__ == "__main__":
    main()

I created a batch generator to yield images and labels in batches. I included print statements to debug and verify the shapes of the images and labels before passing them to the model.

I utilized the train_on_batch method to train the model with batches from the generator. I checked that the shapes of batch_images and batch_labels matched before each training step.