TensorFlow — ValueError: аргумент `output` должен иметь ранг (ndim) `target.ndim - 1`

Я пытаюсь обучить последовательную модель на наборе данных TinyImageNet, используя TensorFlow.

TinyImageNet имеет 200 классов, 100 тысяч строк для обучения и 10 тысяч для проверки. Он имеет два столбца: image и label.

Структура экземпляра:

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

Я предварительно обрабатываю данные, чтобы стандартизировать значения каналов RGB от 0 до 1, как показано здесь.

Однако при подгонке данных к модели я получаю следующую ошибку:

---------------------------------------------------------------------------
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)

Полный код:

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

Я попытался изменить параметр input_shapelayers, а также изменить параметры shapeTensorSpecoutput_signature на None или просто пусто, как я видел в некоторых ответах StackOverflow, но это дало TypeError: 'generator' yielded an element of shape (64, 64, 3) where an element of shape () was expected.

Я также пробовал использовать ImageDataGenerator(rescale=1./255) с такими генераторами, как:

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"
)

но это дало мне TypeError: 'generator' yielded an element of shape (0, 64, 64, 3) where an element of shape (64, 64, 3) was expected. из-за формы генератора x, которая по какой-то причине равна (0, 64, 64, 3):

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

Я использую TensorFlow 2.16.1 и Python 3.11.9 и попробовал перейти на TF 2.14, но ничего не изменилось.

Любая помощь будет принята с благодарностью.