Tensorflow：自定义层输出命名

Question

我正在使用带有字典的 tf.data 管道来训练我的模型，因此我的模型的输入和输出名称必须与我的 tf.data. 数据集中的字典键匹配。要存档，我需要一种方法来控制自定义图层的输出名称。

我不知道该怎么做。请参阅以下示例：

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, name=None):
        super(CustomLayer, self).__init__(name=name)
        self.dense = tf.keras.layers.Dense(32)

    def __call__(self, inputs):
        x = self.dense(inputs)
        x = tf.add(x, 42)
        return x

m = tf.keras.models.Sequential()
m.add(tf.keras.Input(shape=(100,)))
m.add(tf.keras.layers.Dense(84))
m.add(CustomLayer(name='custom_layer'))

m.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 84)                8484      
_________________________________________________________________
dense_1 (Dense)              (None, 32)                2720      
_________________________________________________________________
tf.math.add (TFOpLambda)     (None, 32)                0         
=================================================================
Total params: 11,204
Trainable params: 11,204
Non-trainable params: 0

我期望或想要实现的是像下面这样的图层命名：

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 84)                8484      
_________________________________________________________________
custom_layer                 (None, 32)                2720             
=================================================================
Total params: 11,204
Trainable params: 11,204
Non-trainable params: 0

Answer 1

嗯，原来问题是我覆盖了 __call__ 而不是 call。以下解决了上述问题：

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, name):
        super(CustomLayer, self).__init__(name=name)
        self.dense = tf.keras.layers.Dense(32, name='abc')

    def call(self, inputs):
        x = self.dense(inputs)
        x = x * 2
        return x

m = tf.keras.models.Sequential()
m.add(tf.keras.Input(shape=(100,)))
m.add(tf.keras.layers.Dense(84))
m.add(CustomLayer(name='some_name'))

m.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense (Dense)                (None, 84)                8484      
_________________________________________________________________
some_name (CustomLayer)      (None, 32)                2720             
=================================================================
Total params: 11,204
Trainable params: 11,204
Non-trainable params: 0

Tensorflow：自定义层输出命名

Tensorflow: Custom Layer Output Naming

python

tensorflow

tensorflow2.0