TF.Keras 中具有多输出模型的单损失

Question

我使用tensorflow 的数据集，这样y 是一个包含 6 张量的字典，我都在一个看起来像这样的损失函数中使用它们：

def CustomLoss():
    def custom_loss(y_true, y_pred):
        a = tf.keras.losses.binary_crossentropy(y_true['a_0'], y_pred[0]) * y_true['a_1']
        b = tf.square(y_true['b_0'] - y_pred[1]) * y_true['b_1']
        c = tf.abs(y_true['c_0'] - y_pred[2]) * y_true['c_1']
        return a + b + c
    return custom_loss

我有一个具有 3 个不同形状输出的模型。当我编译模型并调用fit 方法时，我得到Value Error

model.compile(optimizer=optimizer, loss=CustomLoss())
model.fit(dataset, epochs=10)

ValueError: Found unexpected keys that do not correspond to any 
Model output: dict_keys(['a_0', 'a_1', 'b_0', 'b_1', 'c_0', 'c_1']). 
Expected: ['output_0', 'output_1', 'output_2']

其中output_0, 'output_1', 'output_2' 是输出层的名称。

我认为通过数据集中的键命名输出层应该可以解决问题，但问题是我在数据集中有 6 个张量并且只有 3 个输出。我知道我可以使用单个数据集地面实况张量为每个输出分配一个损失函数，但我需要再次将至少两个张量作为 GT 传递。

到目前为止，我使用了自定义训练循环，但我更愿意使用 fit 方法。我正在使用tensorflow 2.3.1

编辑：

示例模型：

inputs = x = tf.keras.layers.Input((256, 256, 3))
x = tf.keras.applications.ResNet50(include_top=False, weights=None)(x)
x1 = tf.keras.layers.Flatten()(x)
x1 = tf.keras.layers.Dense(2, name='output_1')(x1)
x2 = tf.keras.layers.Conv2D(256, 1, name='output_2')(x)
x3 = tf.keras.layers.Flatten()(x)
x3 = tf.keras.layers.Dense(64, name='output_3')(x3)
model = tf.keras.Model(inputs=inputs, outputs=[x1, x2, x3])

自定义训练循环：

avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)
for epoch in range(1, epochs+1):
    for batch, (images, labels) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            outputs = model(images, training=False)
            reg_loss = tf.reduce_sum(model.losses)
            pred_loss = loss(labels, outputs)
            total_loss = tf.reduce_sum(pred_loss) + reg_loss
        grads = tape.gradient(total_loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        avg_loss.update_state(total_loss)
    print(f'Epoch {epoch}/{epochs} - Loss: {avg_loss.result().numpy()}')
    avg_loss.reset_states()

最少的可重现代码：

import tensorflow as tf

def CustomLoss():
    def custom_loss(y_true, y_pred):
        a = tf.keras.losses.binary_crossentropy(y_true['a_0'], y_pred[0]) * y_true['a_1']
        b = tf.square(y_true['b_0'] - y_pred[1]) * y_true['b_1']
        b = tf.reduce_sum(b, axis=(1, 2, 3))
        c = tf.abs(y_true['c_0'] - y_pred[2]) * y_true['c_1']
        c = tf.reduce_sum(c, axis=1)
        return a + b + c
    return custom_loss

dataset = tf.data.Dataset.from_tensors((
    tf.random.uniform((256, 256, 3)),
    {'a_0': [0., 1.], 'a_1': [1.], 'b_0': tf.random.uniform((8, 8, 256)), 'b_1': [1.], 'c_0': tf.random.uniform((64,)), 'c_1': [1.]}
))
dataset = dataset.batch(1)

inputs = x = tf.keras.layers.Input((256, 256, 3))
x = tf.keras.applications.ResNet50(include_top=False, weights=None)(x)
x1 = tf.keras.layers.Flatten()(x)
x1 = tf.keras.layers.Dense(2, name='output_1')(x1)
x2 = tf.keras.layers.Conv2D(256, 1, name='output_2')(x)
x3 = tf.keras.layers.Flatten()(x)
x3 = tf.keras.layers.Dense(64, name='output_3')(x3)
model = tf.keras.Model(inputs=inputs, outputs=[x1, x2, x3])

optimizer = tf.keras.optimizers.Adam(1e-3)
model.compile(optimizer=optimizer, loss=CustomLoss())
model.fit(dataset, epochs=1)

Answer 1

这是适合您情况的一种方法。我们仍将使用自定义训练循环，但也会通过自定义此方法来利用方便的.fit 方法。请查看文档以获取更多详细信息：Customizing what happens in fit()

---

这是一个简单的演示，扩展您的可重现代码。

import tensorflow as tf

# data set 
dataset = tf.data.Dataset.from_tensors((
    tf.random.uniform((256, 256, 3)),
    {'a_0': [0., 1.], 'a_1': [1.], 'b_0': tf.random.uniform((8, 8, 256)),
     'b_1': [1.], 'c_0': tf.random.uniform((64,)), 'c_1': [1.]}
))
dataset = dataset.batch(1)

# custom loss 
def loss(y_true, y_pred):
        a = tf.keras.losses.binary_crossentropy(y_true['a_0'], y_pred[0]) * y_true['a_1']
        b = tf.square(y_true['b_0'] - y_pred[1]) * y_true['b_1']
        b = tf.reduce_sum(b, axis=(1, 2, 3))
        c = tf.abs(y_true['c_0'] - y_pred[2]) * y_true['c_1']
        c = tf.reduce_sum(c, axis=1)
        return a + b + c

自定义模型

这基本上覆盖了train_step，它将在每批数据上重复运行。

avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)

class custom_fit(tf.keras.Model):
    def train_step(self, data):
        images, labels = data
        with tf.GradientTape() as tape:
            outputs = self(images, training=True) # forward pass 
            reg_loss = tf.reduce_sum(self.losses)
            pred_loss = loss(labels, outputs)
            total_loss = tf.reduce_sum(pred_loss) + reg_loss
        gradients = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
        avg_loss.update_state(total_loss)
        return {"loss": avg_loss.result()}
    
    @property
    def metrics(self):
        # We list our `Metric` objects here so that `reset_states()` can be
        # called automatically at the start of each epoch
        # or at the start of `evaluate()`.
        # If you don't implement this property, you have to call
        # `reset_states()` yourself at the time of your choosing.
        return [avg_loss]

构建模型

# model 
inputs = x = tf.keras.layers.Input((256, 256, 3))
x = tf.keras.applications.ResNet50(include_top=False, weights=None)(x)
x1 = tf.keras.layers.Flatten()(x)
x1 = tf.keras.layers.Dense(2, name='output_1')(x1)
x2 = tf.keras.layers.Conv2D(256, 1, name='output_2')(x)
x3 = tf.keras.layers.Flatten()(x)
x3 = tf.keras.layers.Dense(64, name='output_3')(x3)

# simply pass input and outps to the custom model
custom_model = custom_fit(inputs=[inputs], 
                          outputs=[x1, x2, x3])

编译和拟合

custom_model.compile(optimizer='adam')
custom_model.fit(dataset, epochs=5, verbose=2)

Epoch 1/5
1/1 - 6s - loss: 73784.0078
Epoch 2/5
1/1 - 1s - loss: 64882.8984
Epoch 3/5
1/1 - 1s - loss: 54760.2500
Epoch 4/5
1/1 - 1s - loss: 47696.7031
Epoch 5/5
1/1 - 1s - loss: 40574.6328