基于MobileNet的犬只图像多任务分类模型训练报错修复咨询

错误原因 #

报错logits and labels must be broadcastable的核心是输入数据与标签的维度完全不匹配：

• 自定义生成器把两个独立数据集的图像拼接成一个batch（比如第一个batch是11张color数据集图 + 6张breed数据集图，共17张），但返回的标签是[color_labels(11个), breed_labels(6个)]。
• 模型的多任务设计是每张输入图像都要输出color和breed两个预测结果，所以模型会对这17张图的每一张都生成color和breed的logits。计算loss时，color分支的logits是17×12，但color_labels只有11×3；breed分支的logits是17×6，breed_labels只有6×3，维度完全对不上，导致报错。
• 另外代码中Dense层的单元数和实际数据集类别数不匹配：color分支设了12，但你的color数据集只有3类；breed分支设了6，但实际只有3类，这也会导致后续维度问题。

解决步骤 #

1. 重构数据集：使用同时标注双任务的单一数据集 #

多任务模型要求每张图像同时具备color和breed两个标签，所以需要把两个数据集合并，确保每个样本都有两个任务的标注。推荐两种结构：

方式1：按联合标签命名子目录

combined_dataset/
    black_poodle/
        img1.jpg
        img2.jpg
    yellow_labrador/
        img3.jpg
        ...

之后在代码中解析目录名，拆分出color和breed标签。

方式2：使用CSV标注文件

创建labels.csv，格式如下：

image_path,color,breed
./images/img1.jpg,black,poodle
./images/img2.jpg,yellow,labrador
...

2. 修改自定义数据生成器 #

以CSV标注方式为例，重写生成器，确保每个batch的每张图像对应两个标签：

import pandas as pd
import os

class MultiTaskDataGenerator(tf.keras.utils.Sequence):
    def __init__(self, csv_path, img_dir, batch_size, target_size, color_classes, breed_classes):
        self.df = pd.read_csv(csv_path)
        self.img_dir = img_dir
        self.batch_size = batch_size
        self.target_size = target_size
        self.color_classes = color_classes
        self.breed_classes = breed_classes
        # 标签映射
        self.color_map = {cls:i for i,cls in enumerate(color_classes)}
        self.breed_map = {cls:i for i,cls in enumerate(breed_classes)}
        self.data_generator = ImageDataGenerator(rescale=1.0/255.0)

    def __len__(self):
        return len(self.df) // self.batch_size

    def __getitem__(self, index):
        # 获取当前batch的样本
        batch_df = self.df.iloc[index*self.batch_size : (index+1)*self.batch_size]
        images = []
        color_labels = []
        breed_labels = []

        for _, row in batch_df.iterrows():
            # 加载并预处理图像
            img_path = os.path.join(self.img_dir, row['image_path'])
            img = tf.keras.preprocessing.image.load_img(img_path, target_size=self.target_size)
            img = tf.keras.preprocessing.image.img_to_array(img)
            img = self.data_generator.standardize(img)
            images.append(img)
            
            # 转换为one-hot标签
            color_label = tf.one_hot(self.color_map[row['color']], len(self.color_classes))
            breed_label = tf.one_hot(self.breed_map[row['breed']], len(self.breed_classes))
            color_labels.append(color_label)
            breed_labels.append(breed_label)

        # 转换为tensor
        images = tf.stack(images)
        color_labels = tf.stack(color_labels)
        breed_labels = tf.stack(breed_labels)
        
        return images, {'color_output': color_labels, 'breed_output': breed_labels}

3. 修正模型的输出层单元数 #

根据实际数据集的类别数调整Dense层：

# 假设实际color有3类，breed有3类
color_branch = Dense(3, activation='softmax', name='color_output')(color_branch)
breed_branch = Dense(3, activation='softmax', name='breed_output')(breed_branch)

4. 训练模型 #

# 初始化生成器
color_classes = ['black', 'yellow', 'brown'] # 替换为你的实际类别
breed_classes = ['poodle', 'labrador', 'golden_retriever'] # 替换为你的实际类别
data_gen = MultiTaskDataGenerator(
    csv_path='./labels.csv',
    img_dir='./images',
    batch_size=32,
    target_size=(224,224),
    color_classes=color_classes,
    breed_classes=breed_classes
)

# 训练
model.fit(data_gen, epochs=10)

额外提示 #

• 如果暂时无法合并数据集，也可以对两个数据集做对齐（确保每个样本在两个数据集中都有对应图像），但这种方式效率低且容易出错，推荐直接使用同时标注的数据集。
• 训练时可以给两个任务设置不同的loss权重，比如loss_weights={'color_output': 0.3, 'breed_output': 0.7}，平衡两个任务的学习优先级。

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