在 TF2 的损失函数中使用中间层的输出

Using the output of intermediate layers in a loss function in TF2

作为我的 TF2 学习的一部分,我正在尝试在 Tensorflow 2 中复制 OpenPose 的训练,但为了能够做到这一点,我需要在我的损失函数中使用 S、L 中间层的输出。

我试过使用函数 API,但我似乎无法从 S/L 层获得输出,无法根据需要在损失函数中使用它们。我可以看到这可能如何通过子类化实现,但这会增加复杂性并且不适合调试。在我学习的这个阶段,调试和易用性可能是一个巨大的优势。

有什么方法可以用函数 API 或顺序模型来做这种类型的模型?

是的,函数式和顺序式 keras 模型都支持这一点。您始终可以传递一个 dict,其中包含图层名称作为键,损失函数作为值。这是演示这一点的代码。

如果您想从头开始构建模型,只需将图层添加为模型的输出之一即可。

import tensorflow as tf

img = tf.keras.Input([128, 128, 3], name='image')
conv_1 = tf.keras.layers.Conv2D(16, 3, 1, name='conv_1')(img)
conv_2 = tf.keras.layers.Conv2D(16, 3, 1, name='conv_2')(conv_1)
conv_3 = tf.keras.layers.Conv2D(16, 3, 1, name='conv_3')(conv_2)
conv_4 = tf.keras.layers.Conv2D(16, 3, 1, name='conv_4')(conv_3)
conv_5 = tf.keras.layers.Conv2D(16, 3, 1, name='conv_5')(conv_4)
avg_pool = tf.keras.layers.GlobalAvgPool2D(name='avg_pool')(conv_5)
output = tf.keras.layers.Dense(1, activation='sigmoid')(avg_pool)

model = tf.keras.Model(inputs=[img], outputs=[output, conv_5])
print(model.outputs)

输出:

[<tf.Tensor 'dense/Sigmoid:0' shape=(None, 1) dtype=float32>,
<tf.Tensor 'conv_5/BiasAdd:0' shape=(None, 118, 118, 16) dtype=float32>]

但是如果您正在使用已经构建的模型,您可以使用 model.get_layer 方法访问图层并创建新模型:

intermediate_layer = model.get_layer('avg_pool').output
new_model = tf.keras.Model(inputs=model.inputs, outputs=model.outputs + [intermediate_layer])
print(new_model.outputs)

输出:


[<tf.Tensor 'dense/Sigmoid:0' shape=(None, 1) dtype=float32>,
<tf.Tensor 'conv_5/BiasAdd:0' shape=(None, 118, 118, 16) dtype=float32>,
<tf.Tensor 'avg_pool/Mean:0' shape=(None, 16) dtype=float32>]

然后编译您的模型并为每个模型的输出指定一个单独的损失。这些损失可以是字符串,如果它们是 Keras 提供的默认损失,或者这些可以是实现损失函数的可调用对象。

new_model.compile(optimizer='sgd',
              loss={
                  'dense': 'binary_crossentropy',
                  'conv_5': 'mse',
                  'avg_pool': 'mae'
              })

一些虚拟数据和标签

images = tf.random.normal([100, 128, 128, 3])
conv_3_labels = tf.random.normal([100, 118, 118, 16])
avg_pool_labels =  tf.random.normal([100, 16])
class_labels = tf.random.uniform([100], 0, 2, tf.float32)

dataset = tf.data.Dataset.from_tensor_slices(
    (images, (class_labels, conv_3_labels, avg_pool_labels))
)
dataset = dataset.batch(4, drop_remainder=True)

培训

new_model.fit(dataset)

输出:

25/25 [==============================] - 2s 79ms/step - loss: 2.4339
- dense_loss: 0.3904 - conv_5_loss: 1.2367 - avg_pool_loss: 0.8068