Keras函数式API中能否为多输出分支配置不同损失函数？

Absolutely—using multiple loss functions for your multi-task network (semantic segmentation + depth estimation) is totally feasible with Keras' Functional API, and it's actually the standard approach for this kind of problem. Let me break down how to implement it, plus some pro tips to make your model train more effectively.

Step 1: Build Your Multi-Output Model #

First, structure your model with a shared backbone (since both tasks benefit from low-level features like edges and textures) and two separate heads for each task. Naming your output layers will make mapping losses later much clearer:

from tensorflow.keras.layers import Input, Conv2D, UpSampling2D
from tensorflow.keras.models import Model
from tensorflow.keras.losses import CategoricalCrossentropy, MeanSquaredError

# Shared backbone (adjust layers to fit your dataset/needs)
inputs = Input(shape=(256, 256, 3))  # Example input shape
x = Conv2D(64, (3,3), activation='relu', padding='same')(inputs)
x = Conv2D(128, (3,3), activation='relu', padding='same', strides=2)(x)
x = Conv2D(256, (3,3), activation='relu', padding='same')(x)

# Semantic segmentation branch
semantic_head = Conv2D(128, (3,3), activation='relu', padding='same')(x)
semantic_head = UpSampling2D(size=(2,2))(semantic_head)
semantic_output = Conv2D(num_classes, (1,1), activation='softmax', name='semantic_output')(semantic_head)

# Depth estimation branch
depth_head = Conv2D(64, (3,3), activation='relu', padding='same')(x)
depth_head = UpSampling2D(size=(2,2))(depth_head)
depth_output = Conv2D(1, (1,1), activation='linear', name='depth_output')(depth_head)

# Create the multi-output model
model = Model(inputs=inputs, outputs=[semantic_output, depth_output])

Step 2: Assign Losses to Each Output #

You have two straightforward ways to assign your desired losses during compilation:

Option 1: Use a List (Order-Matching) #

If you pass a list of losses, Keras will match them to your outputs in the order you defined them (first loss for semantic_output, second for depth_output):

# Define your loss functions
semantic_loss = CategoricalCrossentropy()
depth_loss = MeanSquaredError()

# Compile the model
model.compile(
    optimizer='adam',
    loss=[semantic_loss, depth_loss],
    metrics={
        'semantic_output': ['accuracy'],  # Track segmentation accuracy
        'depth_output': ['mae']           # Track mean absolute error for depth
    }
)

Option 2: Use a Dictionary (Explicit Mapping) #

For clarity (especially if you add more outputs later), use a dictionary to map each named output to its loss:

model.compile(
    optimizer='adam',
    loss={
        'semantic_output': semantic_loss,
        'depth_output': depth_loss
    },
    metrics={
        'semantic_output': ['accuracy'],
        'depth_output': ['mae']
    }
)

Step 3: Training with Multi-Task Labels #

When training, your target data needs to match the output structure. You can pass either a list (matching output order) or a dictionary (matching output names):

List Format #

# Assuming semantic_labels is one-hot encoded (shape: (batch_size, 256, 256, num_classes))
# depth_maps is your continuous depth data (shape: (batch_size, 256, 256, 1))
model.fit(
    x=train_images,
    y=[semantic_labels, depth_maps],
    epochs=50,
    batch_size=16,
    validation_split=0.2
)

Dictionary Format #

model.fit(
    x=train_images,
    y={
        'semantic_output': semantic_labels,
        'depth_output': depth_maps
    },
    epochs=50,
    batch_size=16,
    validation_split=0.2
)

Pro Tips for Better Training #

• Loss Weighting: The two tasks might have loss values of very different magnitudes (e.g., categorical cross-entropy typically ranges 0-1, while MSE could be much larger if depth values are unnormalized). Fix this by adding loss_weights to balance them:

  model.compile(
      optimizer='adam',
      loss={'semantic_output': semantic_loss, 'depth_output': depth_loss},
      loss_weights={'semantic_output': 1.0, 'depth_output': 0.5},  # Adjust weights based on task priority
      metrics={'semantic_output': ['accuracy'], 'depth_output': ['mae']}
  )
  

• Data Normalization: For depth estimation, normalize your depth maps to a 0-1 range (or standardize them) to bring the MSE loss in line with the cross-entropy loss. For semantic segmentation, ensure your labels are one-hot encoded (required for CategoricalCrossentropy).
• Task-Specific Augmentation: Apply augmentations that make sense for both tasks—e.g., horizontal flips work for both, but avoid color jitter if it harms depth estimation.

This setup works perfectly for multi-task learning in Keras, and it’s flexible enough to adjust as you refine your model.

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