我们可以为多输出模型的每个输出定义损失函数。为此，使用模型的最后一层（输出层）的命名。实现这一点的方法之一是通过以下方式。

import tensorflow as tf
from tensorflow.keras import utils
import numpy as np  

(xtrain, ytrain), (_, _) = keras.datasets.mnist.load_data()
y_out_a = utils.to_categorical(ytrain, num_classes=10) 
y_out_b = (ytrain % 2 == 0).astype('float32')
y_out_c = tf.square(tf.cast(ytrain, tf.float32))
batch_size = 32
data_image = tf.data.Dataset.from_tensor_slices(
     xtrain[..., None]
)
data_label = tf.data.Dataset.from_tensor_slices(
     (y_out_a, y_out_b, y_out_c)
)
dataset = tf.data.Dataset.zip((data_image, data_label))
dataset = dataset.shuffle(buffer_size=8 * batch_size)
dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)

x, y = next(iter(dataset))
y[0].shape, y[1].shape, y[2].shape
(TensorShape([32, 10]), TensorShape([32]), TensorShape([32]))

input = keras.Input(shape=(28, 28, 1))
x = layers.Flatten()(input)
x = layers.Dense(128, activation='relu')(x)
out_a = keras.layers.Dense(10, activation='softmax', name='10cls')(x)
out_b = keras.layers.Dense(1, activation='sigmoid', name='2cls')(x)
out_c = keras.layers.Dense(1, activation='linear', name='1rg')(x)
func_model = keras.Model(
    inputs=[input], outputs=[out_a, out_b, out_c]
)

def categorical(y_true, y_pred):
    return keras.losses.CategoricalCrossentropy()(y_true, y_pred) 

def binary(y_true, y_pred):
    return keras.losses.BinaryCrossentropy()(y_true, y_pred) 

def mse(y_true, y_pred):
    return keras.losses.MeanSquaredError()(y_true, y_pred) 

# compile the model with target loss fn
func_model.compile(
    # you can use what you want
    loss = {
        "10cls": categorical,
        # "2cls": binary,
        # "1rg": mse,
    },
    optimizer = keras.optimizers.Adam()
)

func_model.fit(
    dataset.take(100), 
)
4ms/step - loss: 17.5582 - 10cls_loss: 17.5582

一些资源，这可能也会有帮助

• Single loss with Multiple output model in TF.Keras