如何为多输出Keras模型编写接收所有输出的自定义损失函数？

错误根源 #

在TensorFlow 2.8.4的多输出Keras模型中，当使用继承自tf.keras.losses.Loss的自定义损失类时，Keras会将多个输出对应的真实值（y_true）和预测值（y_pred）堆叠成单个张量传入call方法，而非你假设的张量列表。此时尝试解包张量（y_true1, y_true2 = y_true）会触发迭代张量的操作，这在TensorFlow计算图模式下是被禁止的，从而抛出OperatorNotAllowedInGraphError。

解决方案 #

方案1：改用函数式自定义损失（推荐）

函数式损失会直接接收对应多个输出的y_true和y_pred列表，完全匹配你最初的代码逻辑：

import numpy as np
from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
import tensorflow as tf

n_samples = 1000
n_features = 10
X = np.random.random((n_samples, n_features))
Y = [np.random.random((n_samples, 1)),
     np.random.random((n_samples, 1))]

# 定义函数式自定义损失
def custom_loss(y_true, y_pred):
    y_true1, y_true2 = y_true
    y_pred1, y_pred2 = y_pred

    mse1 = tf.reduce_mean(tf.square(y_true1 - y_pred1))
    mse2 = tf.reduce_mean(tf.square(y_true2 - y_pred2))

    return mse1 + mse2

# 构建多输出模型
input_layer = Input(shape=(n_features,))
x = Dense(32, activation='relu')(input_layer)
outputs = [Dense(1)(x), Dense(1)(x)]

model = Model(inputs=input_layer, outputs=outputs)
model.compile(optimizer=Adam(), loss=custom_loss)

model.fit(X, Y, epochs=2, batch_size=32)

方案2：修改Loss类处理堆叠后的张量

如果坚持使用Loss子类，可通过张量切片分离出两个输出对应的真实值和预测值：

import numpy as np
from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
import tensorflow as tf

n_samples = 1000
n_features = 10
X = np.random.random((n_samples, n_features))
Y = [np.random.random((n_samples, 1)),
     np.random.random((n_samples, 1))]

class CustomLoss(tf.keras.losses.Loss):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def call(self, y_true, y_pred):
        # 从堆叠张量中切片获取两个输出的对应值
        y_true1 = y_true[:, 0:1]
        y_true2 = y_true[:, 1:2]
        y_pred1 = y_pred[:, 0:1]
        y_pred2 = y_pred[:, 1:2]

        mse1 = tf.reduce_mean(tf.square(y_true1 - y_pred1))
        mse2 = tf.reduce_mean(tf.square(y_true2 - y_pred2))

        return mse1 + mse2

# 构建多输出模型
input_layer = Input(shape=(n_features,))
x = Dense(32, activation='relu')(input_layer)
outputs = [Dense(1)(x), Dense(1)(x)]

model = Model(inputs=input_layer, outputs=outputs)
model.compile(optimizer=Adam(), loss=CustomLoss())

model.fit(X, Y, epochs=2, batch_size=32)

方案3：为每个输出指定单独的损失函数

若两个输出的损失逻辑独立，可直接为每个输出分配单独的MSE损失，Keras会自动将总损失求和：

import numpy as np
from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import MeanSquaredError

n_samples = 1000
n_features = 10
X = np.random.random((n_samples, n_features))
Y = [np.random.random((n_samples, 1)),
     np.random.random((n_samples, 1))]

# 构建多输出模型
input_layer = Input(shape=(n_features,))
x = Dense(32, activation='relu')(input_layer)
outputs = [Dense(1)(x), Dense(1)(x)]

model = Model(inputs=input_layer, outputs=outputs)
# 为每个输出指定MSE损失，总损失为两者之和
model.compile(optimizer=Adam(), loss=[MeanSquaredError(), MeanSquaredError()])

model.fit(X, Y, epochs=2, batch_size=32)

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