Keras: one-hot for labels in `image_dataset_from_directory`

Question

I am trying to do a binary image classification using efficientNet. The following is my code.

import matplotlib.pyplot as plt
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.preprocessing import image_dataset_from_directory
from tensorflow.keras.models import Sequential
from tensorflow.keras import layers
from tensorflow.data import AUTOTUNE


DATA_DIR = "img/"
IMG_SIZE = 224 
NUM_CLASSES = 2
EPOCH = 50

def train_val_split(DATA_DIR, IMG_SIZE):
    val_data = image_dataset_from_directory(
                    DATA_DIR,
                    labels="inferred",
                    label_mode="binary",
                    color_mode="rgb",
                    batch_size=32,
                    image_size=(IMG_SIZE, IMG_SIZE),
                    validation_split=0.2,
                    subset="training",
                    seed=1
                )

    train_data = image_dataset_from_directory(
                    DATA_DIR,
                    labels="inferred",
                    label_mode="binary",
                    color_mode="rgb",
                    batch_size=32,
                    image_size=(IMG_SIZE, IMG_SIZE),
                    validation_split=0.2,
                    subset="validation",
                    seed=1
                )

    train_data = train_data.cache().prefetch(buffer_size=AUTOTUNE)
    val_data = val_data.cache().prefetch(buffer_size=AUTOTUNE)
    return train_data, val_data


def model_arch(NUM_CLASSES, IMG_SIZE):
    """efficientnet transfer learning"""
    inputs = layers.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
    img_augmentation = Sequential(
        [
            layers.RandomRotation(factor=0.15),
            layers.RandomTranslation(height_factor=0.1, width_factor=0.1),
            layers.RandomFlip(),
            layers.RandomContrast(factor=0.1),
        ],
        name="img_augmentation",
    )

    x = img_augmentation(inputs)
    # model = EfficientNetB0(include_top=False, input_tensor=x, weights="imagenet")
    model = EfficientNetB0(include_top=False, input_tensor=x, weights='model/efficientnetb0_notop.h5')

    # Freeze the pretrained weights
    model.trainable = False

    # Rebuild top
    x = layers.GlobalAveragePooling2D(name="avg_pool")(model.output)
    x = layers.BatchNormalization()(x)

    top_dropout_rate = 0.2
    x = layers.Dropout(top_dropout_rate, name="top_dropout")(x)
    outputs = layers.Dense(NUM_CLASSES, activation="softmax", name="pred")(x)

    # Compile
    model = tf.keras.Model(inputs, outputs, name="EfficientNet")
    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)
    model.compile(
        optimizer=optimizer, 
        loss="binary_crossentropy", 
        metrics=["accuracy"]
    )
    return model


if __name__ == "__main__":
    train_data, val_data = train_val_split(DATA_DIR, IMG_SIZE)
    model = model_arch(NUM_CLASSES, IMG_SIZE)
    hist = model.fit(train_data, 
                     epochs=EPOCH, 
                     validation_data=val_data, 
                     verbose=1)

However, I encountered the following error.

ValueError: logits and labels must have the same shape ((None, 2) vs (None, 1))

I found out that this is because the labels loaded in image_dataset_from_directory is not one-hot encoded.

print(train_data)
<PrefetchDataset shapes: ((None, 224, 224, 3), (None, 1)), types: (tf.float32, tf.float32)>

How can I tweak the code for train_data and val_data so that it can fit into the model without issues?

Thanks.

Answer 1

You need to change label_mode to a better option.

In your case, I think you want label_mode='categorical',.

As @Dr. Snoopy said the information is here: https://www.tensorflow.org/api_docs/python/tf/keras/utils/image_dataset_from_directory

Answer 2

Managed to figure out the answer!

import tensorflow as tf

# one-hot encoding
train_data = train_data.map(lambda x, y: (x, tf.one_hot(y, depth=NUM_CLASSES)))
val_data = val_data.map(lambda x, y: (x, tf.one_hot(y, depth=NUM_CLASSES)))