Keras多任务共享权重网络中可切换尺寸自定义损失函数问题问询

嘿，我完全懂你在开发共享权重多任务分类器时遇到的这个头疼问题——想用K.switch实现任务切换训练，但不同任务的y_pred和y_true维度不匹配，确实很容易卡壳。结合你提到的三个任务场景，我给你整理了几个实用的解决思路：

思路1：给所有任务套统一维度的“容器” #

我们可以先给所有任务的标签和预测结果设置一个统一的外层维度，让它们进入损失函数时维度一致，再通过任务标识提取对应任务的真实数据计算损失。

举个具体例子，假设你的三个任务分别是：

• 任务A：二分类（输出维度(batch_size, 1)）
• 任务B：三分类（输出维度(batch_size, 3)）
• 任务C：五标签多分类（输出维度(batch_size, 5)）

我们可以把所有y_true包装成(batch_size, 5)的张量（取最大维度5），多余位置用0填充；对应的y_pred也统一输出(batch_size, 5)，再用任务ID判断当前要计算的损失分支：

import tensorflow as tf
from tensorflow.keras import backend as K

def multi_task_loss(y_true, y_pred, task_id):
    # task_id是形状为(batch_size,1)的整数张量，0=任务A，1=任务B，2=任务C
    def task_a_loss():
        # 提取任务A的标签和预测值
        true_a = K.expand_dims(y_true[:, 0], axis=-1)
        pred_a = K.expand_dims(y_pred[:, 0], axis=-1)
        return K.binary_crossentropy(true_a, pred_a)
    
    def task_b_loss():
        true_b = y_true[:, 1:4]
        pred_b = y_pred[:, 1:4]
        return K.categorical_crossentropy(true_b, pred_b)
    
    def task_c_loss():
        true_c = y_true[:, :5]
        pred_c = y_pred[:, :5]
        return K.binary_crossentropy(true_c, pred_c)
    
    # 用K.switch根据任务ID选择对应损失计算分支
    loss = K.switch(K.equal(task_id, 0), task_a_loss(),
                   K.switch(K.equal(task_id, 1), task_b_loss(),
                           task_c_loss()))
    return K.mean(loss)

思路2：用自定义生成器动态匹配维度 #

另一种方式是从数据输入阶段解决——自定义数据生成器，每次训练只输出当前任务对应的标签和网络输出维度。比如训练任务A时，生成器只返回任务A的标签，网络也只输出任务A的维度；切换任务时，生成器和网络输出同步切换。

这种方式需要在网络结构里做动态输出调整，比如用Lambda层根据任务ID选择输出分支：

from tensorflow.keras.layers import Input, Dense, Lambda
from tensorflow.keras.models import Model

# 共享特征提取层
shared_input = Input(shape=(input_dim,))
shared_features = Dense(256, activation='relu')(shared_input)
shared_features = Dense(128, activation='relu')(shared_features)

# 任务ID输入
task_id_input = Input(shape=(1,), dtype='int32')

# 各任务输出分支
task_a_output = Dense(1, activation='sigmoid')(shared_features)
task_b_output = Dense(3, activation='softmax')(shared_features)
task_c_output = Dense(5, activation='sigmoid')(shared_features)

# 用Lambda层根据任务ID选择最终输出
def select_output(x):
    features, task_id = x
    return K.switch(K.equal(task_id, 0), task_a_output,
                   K.switch(K.equal(task_id, 1), task_b_output,
                           task_c_output))

final_output = Lambda(select_output)([shared_features, task_id_input])

# 构建模型
model = Model(inputs=[shared_input, task_id_input], outputs=final_output)

训练时，你只需要传入对应任务的标签和任务ID，直接用该任务的标准损失函数即可，完全不会有维度不匹配的问题。

思路3：用掩码机制忽略无关维度 #

如果不想修改输入输出结构，可以给每个任务的损失添加掩码，让无关维度的损失被置为0，不参与梯度计算：

def masked_multi_task_loss(y_true, y_pred, task_id):
    # 生成各任务的掩码（当前任务为1，其余为0）
    mask_a = K.cast(K.equal(task_id, 0), K.floatx())
    mask_b = K.cast(K.equal(task_id, 1), K.floatx())
    mask_c = K.cast(K.equal(task_id, 2), K.floatx())
    
    # 计算各任务损失并加权掩码
    loss_a = K.binary_crossentropy(y_true[:, :1], y_pred[:, :1]) * mask_a
    loss_b = K.categorical_crossentropy(y_true[:, 1:4], y_pred[:, 1:4]) * mask_b
    loss_c = K.binary_crossentropy(y_true[:, :5], y_pred[:, :5]) * mask_c
    
    # 求和取平均（无关任务损失已被掩码置0）
    total_loss = K.mean(loss_a + loss_b + loss_c)
    return total_loss

这种方式的优势是网络输出可以保留所有任务维度的拼接，训练时只需要传入任务ID，掩码会自动过滤掉无关任务的损失计算。

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