基于Keras的多输出CNN优化：分类交叉熵与argmax实现多任务预测

针对同时预测4个0-1连续变量和1个4分类变量的需求，正确的实现方式是用单模型承载两个独立输出头，分别适配分类与回归任务的特性：

• 分类头：输出维度设为4，用softmax激活，搭配CategoricalCrossentropy损失函数，对应独热编码的分类标签
• 回归头：输出维度设为4，用linear激活，搭配MeanSquaredError损失函数，对应4个连续数值变量

以下是具体的代码修改方案与完整实现：

一、核心修改要点 #

1. 数据准备简化：无需为独热编码的每个类别单独创建DataFrame列，直接用原始分类标签生成独热编码数组
2. 模型结构优化：合并分类输出为单个softmax头，回归输出为单个linear头，共享卷积特征提取部分
3. 编译配置适配：为不同输出头指定对应损失函数与评估指标
4. 预测后处理：用np.argmax()将分类头的概率输出转为最终类别编号

二、完整修改后代码 #

from tensorflow.keras.utils import to_categorical
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras import layers
import tensorflow.keras as keras
import pandas as pd
import numpy as np


def baselineCNNModel(train_df, test_df, OUTPUT_DIR_TRAIN, OUTPUT_DIR_TEST, debug=False):
    """
    构建用于多输出预测的基线CNN模型。
    输入为包含道路标志的图像，目标为类别标签和边界框坐标。
    """

    print("\nrunSimpleModel\n")
    train_dataset = train_df[['Class Number', 'Center in X', 'Center in Y', 'Width', 'Height', 'Image Filename']]
    test_dataset = test_df[['Class Number', 'Center in X', 'Center in Y', 'Width', 'Height', 'Image Filename']]

    # 生成分类标签的独热编码，无需添加到DataFrame
    train_class_onehot = to_categorical(train_dataset['Class Number'], num_classes=4)
    test_class_onehot = to_categorical(test_dataset['Class Number'], num_classes=4)

    tDIR, sDIR = OUTPUT_DIR_TRAIN, OUTPUT_DIR_TEST
    BS, image_size = 64, (128, 128)

    # 自定义生成器，拆分分类与回归目标
    def multi_output_generator(generator, class_onehot):
        while True:
            x_batch, y_batch = next(generator)
            # 提取分类标签与回归变量
            batch_start = generator.batch_index * BS
            batch_end = batch_start + BS
            class_labels = class_onehot[batch_start:batch_end]
            reg_labels = y_batch[:, 1:]  # y_batch第一列是原始分类编号，后4列是回归变量
            yield x_batch, {'class_output': class_labels, 'bbox_output': reg_labels}

    # 数据预处理与增强
    datagen = ImageDataGenerator(
        rescale=1.0 / 255.0,
        validation_split=0.2
    )
    # 基础生成器读取全部目标列
    base_train_generator = datagen.flow_from_dataframe(
        dataframe=train_dataset,
        directory=tDIR,
        x_col="Image Filename",
        y_col=["Class Number", "Center in X", "Center in Y", "Width", "Height"],
        target_size=image_size,
        batch_size=BS,
        class_mode='multi_output',
        subset='training'
    )
    train_generator = multi_output_generator(base_train_generator, train_class_onehot)

    base_val_generator = datagen.flow_from_dataframe(
        dataframe=train_dataset,
        directory=tDIR,
        x_col="Image Filename",
        y_col=["Class Number", "Center in X", "Center in Y", "Width", "Height"],
        target_size=image_size,
        batch_size=BS,
        class_mode='multi_output',
        subset='validation'
    )
    validation_generator = multi_output_generator(base_val_generator, train_class_onehot)

    # 定义CNN模型结构
    input_layer = layers.Input(shape=(image_size[0], image_size[1], 3))
    x = layers.Conv2D(128, (4, 4), activation='relu')(input_layer)
    x = layers.MaxPooling2D((4, 4))(x)
    x = layers.Conv2D(64, (3, 3), activation='relu')(x)
    x = layers.MaxPooling2D((2, 2))(x)
    x = layers.Flatten()(x)
    x = layers.Dense(128, activation='relu')(x)

    # 分类输出头：4类别softmax
    class_output = layers.Dense(4, activation="softmax", name='class_output')(x)
    # 回归输出头：4个连续变量linear激活
    bbox_output = layers.Dense(4, activation="linear", name='bbox_output')(x)

    # 创建多输出模型
    model = keras.Model(inputs=input_layer, outputs=[class_output, bbox_output])

    # 编译模型：指定不同输出的损失与指标
    model.compile(
        optimizer='adam',
        loss={
            'class_output': keras.losses.CategoricalCrossentropy(),
            'bbox_output': keras.losses.MeanSquaredError()
        },
        metrics={
            'class_output': 'accuracy',
            'bbox_output': 'mae'
        }
    )

    # 训练模型
    epochs = 10
    train_steps = len(base_train_generator)
    val_steps = len(base_val_generator)
    history = model.fit(
        train_generator,
        steps_per_epoch=train_steps,
        epochs=epochs,
        validation_data=validation_generator,
        validation_steps=val_steps
    )

    # 模型评估
    evaluation = model.evaluate(validation_generator, steps=val_steps)
    print("\nEvaluation Results:", evaluation)

    # 测试集预测
    test_datagen = ImageDataGenerator(rescale=1.0/255.0)
    base_test_generator = test_datagen.flow_from_dataframe(
        dataframe=test_dataset,
        directory=sDIR,
        x_col="Image Filename",
        y_col=["Class Number", "Center in X", "Center in Y", "Width", "Height"],
        target_size=image_size,
        batch_size=BS,
        class_mode='multi_output',
        shuffle=False  # 保持预测顺序与原始数据一致
    )
    test_generator = multi_output_generator(base_test_generator, test_class_onehot)
    test_steps = len(base_test_generator)

    predictions = model.predict(test_generator, steps=test_steps)
    class_pred_probs, bbox_preds = predictions

    # 将分类概率转为类别编号
    class_preds = np.argmax(class_pred_probs, axis=1)

    # 生成预测结果DataFrame
    prediction_df = pd.DataFrame({
        "Predicted Class Number": class_preds,
        "Center in X": bbox_preds[:, 0],
        "Center in Y": bbox_preds[:, 1],
        "Width": bbox_preds[:, 2],
        "Height": bbox_preds[:, 3],
        'Image Filename': test_dataset['Image Filename'].values
    })

    print("\nPredictions: ")
    print(prediction_df)

    return model, prediction_df

三、关键细节说明 #

• 生成器处理：用multi_output模式读取目标列，自定义生成器拆分分类与回归数据，避免手动维护独热编码列
• 模型结构：共享卷积层提取图像特征，分类与回归头各自独立，适配不同任务的输出要求
• 损失配置：分类任务用交叉熵优化概率分布，回归任务用均方误差优化连续变量，保证两个任务的训练效率
• 预测处理：通过np.argmax()将分类头的4维概率输出转为具体类别编号，直接获取回归头的连续数值结果

内容的提问来源于stack exchange，提问作者Dkasi