TensorFlow - ValueError: Argument `output` must have rank (ndim) `target.ndim - 1`

I'm trying to train a sequential model on the TinyImageNet dataset, using TensorFlow.

TinyImageNet has 200 classes, 100k rows for training and 10k for validating. It has two columns, image and label.

The structure of an instance:

{
  'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=64x64 at 0x1A800E8E190,
  'label': 15
}

I'm preprocessing the data to standardize the RGB channel values between 0 and 1 as seen here.

However, when fitting the data to the model, I'm getting this error :

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[47], line 1
----> 1 model.fit(train_dataset_tf, epochs=10)

File .venv\Lib\site-packages\keras\src\utils\traceback_utils.py:122, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    119     filtered_tb = _process_traceback_frames(e.__traceback__)
    120     # To get the full stack trace, call:
    121     # `keras.config.disable_traceback_filtering()`
--> 122     raise e.with_traceback(filtered_tb) from None
    123 finally:
    124     del filtered_tb

File .venv\Lib\site-packages\keras\src\backend\tensorflow\nn.py:619, in sparse_categorical_crossentropy(target, output, from_logits, axis)
    613     raise ValueError(
    614         "Argument `output` must be at least rank 1. "
    615         "Received: "
    616         f"output.shape={output.shape}"
    617     )
    618 if len(target.shape) != len(output.shape[:-1]):
--> 619     raise ValueError(
    620         "Argument `output` must have rank (ndim) `target.ndim - 1`. "
    621         "Received: "
    622         f"target.shape={target.shape}, output.shape={output.shape}"
    623     )
    624 for e1, e2 in zip(target.shape, output.shape[:-1]):
    625     if e1 is not None and e2 is not None and e1 != e2:

ValueError: Argument `output` must have rank (ndim) `target.ndim - 1`. Received: target.shape=(None, 64, 64, 3), output.shape=(None, 200)

Complete code :

import tensorflow as tf
from tensorflow.keras import layers
from datasets import load_dataset
import numpy as np

train_dataset = load_dataset("Maysee/tiny-imagenet", split="train")
test_dataset = load_dataset("Maysee/tiny-imagenet", split="valid")

def preprocess_image(image, label):
    image = tf.io.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [64, 64])
    image = image / 255.0
    return image, label

def generator(dataset):
    for example in dataset:
        yield example['image'], example['label']

train_dataset_tf = tf.data.Dataset.from_generator(
    lambda: generator(train_dataset),
    output_signature=(
        tf.TensorSpec(shape=(), dtype=tf.string),
        tf.TensorSpec(shape=(64, 64, 3), dtype=tf.int64),
    ),
)

test_dataset_tf = tf.data.Dataset.from_generator(
    lambda: generator(test_dataset),
    output_signature=(
        tf.TensorSpec(shape=(), dtype=tf.string),
        tf.TensorSpec(shape=(64, 64, 3), dtype=tf.int64),
    ),
)

train_dataset_tf = (
    train_dataset_tf.map(preprocess_image).batch(32).prefetch(tf.data.AUTOTUNE)
)
test_dataset_tf = (
    test_dataset_tf.map(preprocess_image).batch(32).prefetch(tf.data.AUTOTUNE)
)

model = tf.keras.Sequential(
    [
        layers.Conv2D(32, (3, 3), activation="relu", input_shape=(64, 64, 3)),
        layers.MaxPooling2D((2, 2)),
        layers.Flatten(),
        layers.Dense(128, activation="relu"),
        layers.Dense(200, activation="softmax"),
    ]
)

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

model.fit(train_dataset_tf, epochs=10) # Error is thrown here

I tried changing the input_shape parameter of the layers, as well as changing the shape parameters of the TensorSpec output_signature to None or just empty as I saw on some StackOverflow answers, but that gave a TypeError: 'generator' yielded an element of shape (64, 64, 3) where an element of shape () was expected.

I also tried using ImageDataGenerator(rescale=1./255) with generators such as :

train_df = pd.DataFrame(train_dataset, dtype=str)
train_generator = train_datagen.flow_from_dataframe(
    train_df,
    directory=None,
    x_col="image",
    y_col="label",
    target_size=(64, 64),
    batch_size=32,
    class_mode="categorical"
)

but this gave me a TypeError: 'generator' yielded an element of shape (0, 64, 64, 3) where an element of shape (64, 64, 3) was expected. because of the generator x shape which is (0, 64, 64, 3) for some reason :

for x, y in train_generator:
    print(x.shape, y.shape) # outputs (0, 64, 64, 3) (0, 0)

I'm using TensorFlow 2.16.1 and Python 3.11.9 and tried downgrading to TF 2.14 but nothing changed.

Any help would be greatly appreciated.

That error is due to some images of the dataset being in only 2 dimensions and not 3 (grayscale for example).

You need to modify your generator to check for such a case. You can handle it like this:

for element in dataset:
   image = np.array(element["img"])
   label = element["label"]

   # duplicate on 3 dims if image grayscale
   if image.ndim == 2:
       image = np.stack([image] * 3, axis=-1)

   yield image, label