ValueError when adding tf.keras.metrics.Precision to model metrics in TensorFlow

I'm trying to apply transfer learning using MobileNet on a custom dataset. My code works fine until I add tf.keras.metrics.Precision(name="precision") to the model's metrics. After adding this metric, I encounter the following error during model.fit:

update_confusion_matrix_variables y_pred.shape.assert_is_compatible_with(y_true.shape) ValueError: Shapes (None, 4) and (None, 1) are incompatible

The relevant code block:

import tensorflow as tf
from tensorflow.keras.utils import image_dataset_from_directory as get_dataset

img_size = (224, 224)
preprocessing_layer = tf.keras.applications.mobilenet.preprocess_input
img_shape = img_size + (3,)

base_model = tf.keras.applications.MobileNet(
    input_shape=img_shape,
    include_top=False,
    weights='imagenet',
)

batch_size = 32
train_set = get_dataset(
    'some path',
    shuffle=True,
    batch_size=batch_size,
    image_size=img_size,
)

class_names = train_set.class_names
num_classes = len(class_names)

base_model.trainable = False
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()

inputs = tf.keras.Input(shape=img_shape)
k = preprocessing_layer(inputs)
k = base_model(k, training=False)
k = global_average_layer(k)
k = tf.keras.layers.Dropout(0.2)(k)

prediction_layer = tf.keras.layers.Dense(
    num_classes,
    activation="softmax"
)
outputs = prediction_layer(k)

model = tf.keras.Model(inputs, outputs)

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=[
        tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
        tf.keras.metrics.Precision(name="precision"),
    ]
)

AUTOTUNE = tf.data.AUTOTUNE

train_set = train_set.prefetch(buffer_size=AUTOTUNE)

epochs = 1
history = model.fit(
    train_set,
    epochs=epochs,
)

The dataset is organized into subdirectories where each subdirectory name is the class label. There are four class labels.

The dataset loading and preprocessing seem to work correctly since the model trains without the tf.keras.metrics.Precision metric.

The issue arises specifically when adding the precision metric.

epochs is intentionally set to 1 for experimental purposes.

Solution

The problem is that you have sparse class labels and Precision can't work with that. The easiest fix would be to one-hot encode your labels and change your sparse loss and metric to the non-sparse versions.

You can change

train_set = get_dataset(
    'some path',
    shuffle=True,
    batch_size=batch_size,
    image_size=img_size,
)
""" [more code from you ...] """

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=[
        tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
        tf.keras.metrics.Precision(name="precision"),
    ]
)

train_set = get_dataset(
    'some path',
    label_mode='categorical',
    shuffle=True,
    batch_size=batch_size,
    image_size=img_size,
)
""" [more code from you ...] """

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
    loss=tf.keras.losses.CategoricalCrossentropy(),
    metrics=[
        tf.keras.metrics.CategoricalAccuracy(name='accuracy'),
        tf.keras.metrics.Precision(name="precision"),
    ]
)

With label_mode='categorical' you get one-hot encoded labels, which work with most of the metrics from keras. You have to change the loss and accuracy to the non-sparse version as shown in the code example above.