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

时间:2020-07-02 16:55:59

标签: tensorflow machine-learning tensorflow2.0

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

我尝试使用功能性API,但似乎无法从S / L层获取输出,从而无法根据需要在损失函数中使用它们。我可以看到子类化如何实现,但这会增加复杂性,而不是调试的理想选择。在我所学的这个阶段,调试和易于使用可能是一大优势。

我可以通过功能性API或顺序模型来进行此类模型吗?

enter image description here

1 个答案:

答案 0 :(得分:2)

是的,功能性和顺序性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