我正在尝试用efficientNet做一个二值图像分类。下面是我的代码。

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)

但是，我遇到了以下错误。

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

我发现这是因为image_dataset_from_directory中加载的标签不是one-hot编码的。

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

如何调整train_data和val_data的代码，使其能够毫无问题地适合模型？
谢谢。