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.